pangea-sdk 6.2.0b1__py3-none-any.whl → 6.3.0__py3-none-any.whl

pangea/services/ai_guard.py

@@ -1,6 +1,10 @@
 from __future__ import annotations
 
-from typing import Any, Dict, Generic, List, Literal, Optional, TypeVar, overload
+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
 from pangea.response import APIRequestModel, APIResponseModel, PangeaResponse, PangeaResponseResult
@@ -17,6 +21,13 @@ MaliciousEntityAction = Literal["report", "defang", "disabled", "block"]
 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
@@ -24,14 +35,14 @@ class CodeDetectionOverride(APIRequestModel):
 
 class LanguageDetectionOverride(APIRequestModel):
     disabled: Optional[bool] = None
-    allow: Optional[List[str]] = None
-    block: Optional[List[str]] = None
-    report: Optional[List[str]] = None
+    allow: Optional[list[str]] = None
+    block: Optional[list[str]] = None
+    report: Optional[list[str]] = None
 
 
 class TopicDetectionOverride(APIRequestModel):
     disabled: Optional[bool] = None
-    block: Optional[List[str]] = None
+    block: Optional[list[str]] = None
 
 
 class PromptInjectionOverride(APIRequestModel):
@@ -179,7 +190,7 @@ class PromptInjectionResult(APIResponseModel):
     action: str
     """The action taken by this Detector"""
 
-    analyzer_responses: List[AnalyzerResponse]
+    analyzer_responses: list[AnalyzerResponse]
     """Triggered prompt injection analyzers."""
 
 
@@ -192,20 +203,20 @@ class PiiEntity(APIResponseModel):
 
 
 class PiiEntityResult(APIResponseModel):
-    entities: List[PiiEntity]
+    entities: list[PiiEntity]
     """Detected redaction rules."""
 
 
 class MaliciousEntity(APIResponseModel):
     type: str
     value: str
-    action: str
+    action: Optional[str] = None
     start_pos: Optional[int] = None
-    raw: Optional[Dict[str, Any]] = None
+    raw: Optional[dict[str, object]] = None
 
 
 class MaliciousEntityResult(APIResponseModel):
-    entities: List[MaliciousEntity]
+    entities: list[MaliciousEntity]
     """Detected harmful items."""
 
 
@@ -215,11 +226,11 @@ class CustomEntity(APIResponseModel):
     action: str
     """The action taken on this Entity"""
     start_pos: Optional[int] = None
-    raw: Optional[Dict[str, Any]] = None
+    raw: Optional[dict[str, object]] = None
 
 
 class CustomEntityResult(APIResponseModel):
-    entities: List[CustomEntity]
+    entities: list[CustomEntity]
     """Detected redaction rules."""
 
 
@@ -233,7 +244,7 @@ class SecretsEntity(APIResponseModel):
 
 
 class SecretsEntityResult(APIResponseModel):
-    entities: List[SecretsEntity]
+    entities: list[SecretsEntity]
     """Detected redaction rules."""
 
 
@@ -266,22 +277,22 @@ class TextGuardDetectors(APIResponseModel):
     prompt_injection: Optional[TextGuardDetector[PromptInjectionResult]] = None
     pii_entity: Optional[TextGuardDetector[PiiEntityResult]] = None
     malicious_entity: Optional[TextGuardDetector[MaliciousEntityResult]] = None
-    custom_entity: Optional[TextGuardDetector[Any]] = None
+    custom_entity: Optional[TextGuardDetector[object]] = None
     secrets_detection: Optional[TextGuardDetector[SecretsEntityResult]] = None
-    profanity_and_toxicity: Optional[TextGuardDetector[Any]] = None
+    profanity_and_toxicity: Optional[TextGuardDetector[object]] = None
     language_detection: Optional[TextGuardDetector[LanguageDetectionResult]] = None
     topic_detection: Optional[TextGuardDetector[TopicDetectionResult]] = None
     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
@@ -345,7 +356,7 @@ class AIGuard(ServiceBase):
         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
 
@@ -373,12 +384,12 @@ class AIGuard(ServiceBase):
     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
 
@@ -400,19 +411,19 @@ class AIGuard(ServiceBase):
             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
 

pangea/services/audit/audit.py

@@ -1,20 +1,21 @@
 # Copyright 2022 Pangea Cyber Corporation
 # Author: Pangea Cyber Corporation
+
+# TODO: Modernize.
+# ruff: noqa: UP006, UP035
+
 from __future__ import annotations
 
 import datetime
 import json
-from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple, Union, cast, overload
-
-from pydantic import TypeAdapter
-from typing_extensions import Literal
+from collections.abc import Mapping
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple, Union
 
 import pangea.exceptions as pexc
 from pangea.config import PangeaConfig
 from pangea.response import PangeaResponse, PangeaResponseResult
 from pangea.services.audit.exceptions import AuditException, EventCorruption
 from pangea.services.audit.models import (
-    AuditSchema,
     DownloadFormat,
     DownloadRequest,
     DownloadResult,
@@ -22,7 +23,6 @@ from pangea.services.audit.models import (
     EventEnvelope,
     EventVerification,
     ExportRequest,
-    ForwardingConfiguration,
     LogBulkRequest,
     LogBulkResult,
     LogEvent,
@@ -40,9 +40,6 @@ from pangea.services.audit.models import (
     SearchRequest,
     SearchResultOutput,
     SearchResultRequest,
-    ServiceConfig,
-    ServiceConfigFilter,
-    ServiceConfigListResult,
 )
 from pangea.services.audit.signing import Signer, Verifier
 from pangea.services.audit.util import (
@@ -63,7 +60,7 @@ from pangea.utils import canonicalize_nested_json
 
 class AuditBase:
     def __init__(
-        self, private_key_file: str = "", public_key_info: dict[str, str] = {}, tenant_id: str | None = None
+        self, private_key_file: str = "", public_key_info: Mapping[str, str] = {}, tenant_id: str | None = None
     ) -> None:
         self.pub_roots: Dict[int, PublishedRoot] = {}
         self.buffer_data: Optional[str] = None
@@ -98,7 +95,7 @@ class AuditBase:
         return input  # type: ignore[return-value]
 
     def _process_log(self, event: dict, sign_local: bool) -> LogEvent:
-        if event.get("tenant_id", None) is None and self.tenant_id:
+        if event.get("tenant_id") is None and self.tenant_id:
             event["tenant_id"] = self.tenant_id
 
         event = {k: v for k, v in event.items() if v is not None}
@@ -229,10 +226,7 @@ class AuditBase:
         tree_sizes.add(result.root.size)
         tree_sizes.difference_update(self.pub_roots.keys())
 
-        if tree_sizes:
-            arweave_roots = get_arweave_published_roots(result.root.tree_name, tree_sizes)
-        else:
-            arweave_roots = {}
+        arweave_roots = get_arweave_published_roots(result.root.tree_name, tree_sizes) if tree_sizes else {}
 
         return tree_sizes, arweave_roots
 
@@ -344,6 +338,7 @@ class AuditBase:
         public_key = get_public_key(audit_envelope.public_key)
 
         if audit_envelope and audit_envelope.signature and public_key:
+            assert audit_envelope.event
             v = Verifier()
             verification = v.verify_signature(
                 audit_envelope.signature,
@@ -393,7 +388,7 @@ class Audit(ServiceBase, AuditBase):
         token: str,
         config: PangeaConfig | None = None,
         private_key_file: str = "",
-        public_key_info: dict[str, str] = {},
+        public_key_info: Mapping[str, str] = {},
         tenant_id: str | None = None,
         logger_name: str = "pangea",
         config_id: str | None = None,
@@ -465,7 +460,7 @@ class Audit(ServiceBase, AuditBase):
             A PangeaResponse where the hash of event data and optional verbose
                 results are returned in the response.result field.
                 Available response fields can be found in our
-                [API documentation](https://pangea.cloud/docs/api/audit#/v1/log).
+                [API documentation](https://pangea.cloud/docs/api/audit#/v1/log-post).
 
         Examples:
             log_response = audit.log(
@@ -513,7 +508,7 @@ class Audit(ServiceBase, AuditBase):
         Returns:
             A PangeaResponse where the hash of event data and optional verbose
                 results are returned in the response.result field.
-                Available response fields can be found in our [API documentation](https://pangea.cloud/docs/api/audit#/v1/log).
+                Available response fields can be found in our [API documentation](https://pangea.cloud/docs/api/audit#/v1/log-post).
 
         Examples:
             response = audit.log_event({"message": "hello world"}, verbose=True)
@@ -551,7 +546,7 @@ class Audit(ServiceBase, AuditBase):
         Returns:
             A PangeaResponse where the hash of event data and optional verbose
                 results are returned in the response.result field.
-                Available response fields can be found in our [API documentation](https://pangea.cloud/docs/api/audit#/v2/log).
+                Available response fields can be found in our [API documentation](https://pangea.cloud/docs/api/audit#/v2/log-post).
 
         Examples:
             log_response = audit.log_bulk(
@@ -592,7 +587,7 @@ class Audit(ServiceBase, AuditBase):
         Returns:
             A PangeaResponse where the hash of event data and optional verbose
                 results are returned in the response.result field.
-                Available response fields can be found in our [API documentation](https://pangea.cloud/docs/api/audit#/v2/log_async).
+                Available response fields can be found in our [API documentation](https://pangea.cloud/docs/api/audit#/v2/log_async-post).
 
         Examples:
             log_response = audit.log_bulk_async(
@@ -667,8 +662,8 @@ class Audit(ServiceBase, AuditBase):
 
         Returns:
             A PangeaResponse[SearchOutput] where the first page of matched events is returned in the
-                response.result field. Available response fields can be found in our [API documentation](https://pangea.cloud/docs/api/audit#/v1/search).
-                Pagination can be found in the [search results endpoint](https://pangea.cloud/docs/api/audit#/v1/results).
+                response.result field. Available response fields can be found in our [API documentation](https://pangea.cloud/docs/api/audit#/v1/search-post).
+                Pagination can be found in the [search results endpoint](https://pangea.cloud/docs/api/audit#/v1/results-post).
 
         Examples:
             response = audit.search(
@@ -927,298 +922,6 @@ class Audit(ServiceBase, AuditBase):
         )
         return self.request.post("v1/download_results", DownloadResult, data=input.model_dump(exclude_none=True))
 
-    def get_service_config(self, config_id: str) -> PangeaResponse[ServiceConfig]:
-        """
-        Get a service config.
-
-        OperationId: audit_post_v1beta_config
-
-        Args:
-            id: The config ID
-        """
-
-        response = self.request.post("v1beta/config", PangeaResponseResult, data={"id": config_id})
-        response.result = TypeAdapter(ServiceConfig).validate_python(response.json["result"])
-        return cast(PangeaResponse[ServiceConfig], response)
-
-    @overload
-    def create_service_config(
-        self,
-        version: Literal[1],
-        name: str,
-        *,
-        cold_query_result_retention: str | None = None,
-        hot_storage: str | None = None,
-        query_result_retention: str | None = None,
-        redact_service_config_id: str | None = None,
-        redaction_fields: Sequence[str] | None = None,
-        retention: str | None = None,
-        vault_key_id: str | None = None,
-        vault_service_config_id: str | None = None,
-        vault_sign: bool | None = None,
-    ) -> PangeaResponse[ServiceConfig]:
-        """
-        Create a v1 service config.
-
-        OperationId: audit_post_v1beta_config_create
-
-        Args:
-            name: Configuration name
-            cold_query_result_retention: Retention window for cold query result / state information.
-            hot_storage: Retention window to keep audit logs in hot storage.
-            query_result_retention: Length of time to preserve server-side query result caching.
-            redact_service_config_id: A redact service config that will be used to redact PII from logs.
-            redaction_fields: Fields to perform redaction against.
-            retention: Retention window to store audit logs.
-            vault_key_id: ID of the Vault key used for signing. If missing, use a default Audit key.
-            vault_service_config_id: A vault service config that will be used to sign logs.
-            vault_sign: Enable/disable event signing.
-        """
-
-    @overload
-    def create_service_config(
-        self,
-        version: Literal[2],
-        name: str,
-        *,
-        schema: AuditSchema,
-        cold_query_result_retention: str | None = None,
-        forwarding_configuration: ForwardingConfiguration | None = None,
-        hot_storage: str | None = None,
-        query_result_retention: str | None = None,
-        redact_service_config_id: str | None = None,
-        retention: str | None = None,
-        vault_key_id: str | None = None,
-        vault_service_config_id: str | None = None,
-        vault_sign: bool | None = None,
-    ) -> PangeaResponse[ServiceConfig]:
-        """
-        Create a v2 service config.
-
-        OperationId: audit_post_v1beta_config_create
-
-        Args:
-            name: Configuration name
-            schema: Audit log field configuration. Only settable at create time.
-            cold_query_result_retention: Retention window for cold query result / state information.
-            forwarding_configuration: Configuration for forwarding audit logs to external systems.
-            hot_storage: Retention window to keep audit logs in hot storage.
-            query_result_retention: Length of time to preserve server-side query result caching.
-            redact_service_config_id: A redact service config that will be used to redact PII from logs.
-            retention: Retention window to store audit logs.
-            vault_key_id: ID of the Vault key used for signing. If missing, use a default Audit key.
-            vault_service_config_id: A vault service config that will be used to sign logs.
-            vault_sign: Enable/disable event signing.
-        """
-
-    @overload
-    def create_service_config(
-        self,
-        version: Literal[3],
-        name: str,
-        *,
-        schema: AuditSchema,
-        cold_storage: str | None = None,
-        hot_storage: str | None = None,
-        warm_storage: str | None = None,
-        redact_service_config_id: str | None = None,
-        vault_service_config_id: str | None = None,
-        vault_key_id: str | None = None,
-        vault_sign: bool | None = None,
-        forwarding_configuration: ForwardingConfiguration | None = None,
-    ) -> PangeaResponse[ServiceConfig]:
-        """
-        Create a v3 service config.
-
-        OperationId: audit_post_v1beta_config_create
-
-        Args:
-            name: Configuration name
-            schema: Audit log field configuration. Only settable at create time.
-            cold_storage: Retention window for logs in cold storage. Deleted afterwards.
-            hot_storage: Retention window for logs in hot storage. Migrated to warm, cold, or deleted afterwards.
-            warm_storage: Retention window for logs in warm storage. Migrated to cold or deleted afterwards.
-            redact_service_config_id: A redact service config that will be used to redact PII from logs.
-            vault_service_config_id: A vault service config that will be used to sign logs.
-            vault_key_id: ID of the Vault key used for signing. If missing, use a default Audit key.
-            vault_sign: Enable/disable event signing.
-            forwarding_configuration: Configuration for forwarding audit logs to external systems.
-        """
-
-    def create_service_config(
-        self,
-        version: Literal[1, 2, 3],
-        name: str,
-        *,
-        cold_query_result_retention: str | None = None,
-        cold_storage: str | None = None,
-        forwarding_configuration: ForwardingConfiguration | None = None,
-        hot_storage: str | None = None,
-        query_result_retention: str | None = None,
-        redact_service_config_id: str | None = None,
-        redaction_fields: Sequence[str] | None = None,
-        retention: str | None = None,
-        schema: AuditSchema | None = None,
-        vault_key_id: str | None = None,
-        vault_service_config_id: str | None = None,
-        vault_sign: bool | None = None,
-        warm_storage: str | None = None,
-    ) -> PangeaResponse[ServiceConfig]:
-        """
-        Create a service config.
-
-        OperationId: audit_post_v1beta_config_create
-
-        Args:
-            name: Configuration name
-            cold_query_result_retention: Retention window for cold query result / state information.
-            cold_storage: Retention window for logs in cold storage. Deleted afterwards.
-            forwarding_configuration: Configuration for forwarding audit logs to external systems.
-            hot_storage: Retention window to keep audit logs in hot storage.
-            query_result_retention: Length of time to preserve server-side query result caching.
-            redact_service_config_id: A redact service config that will be used to redact PII from logs.
-            redaction_fields: Fields to perform redaction against.
-            retention: Retention window to store audit logs.
-            schema: Audit log field configuration. Only settable at create time.
-            vault_key_id: ID of the Vault key used for signing. If missing, use a default Audit key.
-            vault_service_config_id: A vault service config that will be used to sign logs.
-            vault_sign: Enable/disable event signing.
-            warm_storage: Retention window for logs in warm storage. Migrated to cold or deleted afterwards.
-        """
-
-        response = self.request.post(
-            "v1beta/config/create",
-            PangeaResponseResult,
-            data={
-                "cold_query_result_retention": cold_query_result_retention,
-                "cold_storage": cold_storage,
-                "forwarding_configuration": forwarding_configuration,
-                "hot_storage": hot_storage,
-                "name": name,
-                "query_result_retention": query_result_retention,
-                "redact_service_config_id": redact_service_config_id,
-                "redaction_fields": redaction_fields,
-                "retention": retention,
-                "schema": schema,
-                "vault_key_id": vault_key_id,
-                "vault_service_config_id": vault_service_config_id,
-                "vault_sign": vault_sign,
-                "warm_storage": warm_storage,
-                "version": version,
-            },
-        )
-        response.result = TypeAdapter(ServiceConfig).validate_python(response.json["result"])
-        return cast(PangeaResponse[ServiceConfig], response)
-
-    def update_service_config(
-        self,
-        config_id: str,
-        *,
-        name: str,
-        updated_at: datetime.datetime,
-        # Optionals.
-        cold_query_result_retention: str | None = None,
-        cold_storage: str | None = None,
-        forwarding_configuration: ForwardingConfiguration | None = None,
-        hot_storage: str | None = None,
-        query_result_retention: str | None = None,
-        redact_service_config_id: str | None = None,
-        retention: str | None = None,
-        schema: AuditSchema | None = None,
-        vault_key_id: str | None = None,
-        vault_service_config_id: str | None = None,
-        vault_sign: bool | None = None,
-        warm_storage: str | None = None,
-    ) -> PangeaResponse[ServiceConfig]:
-        """
-        Update a service config.
-
-        OperationId: audit_post_v1beta_config_update
-
-        Args:
-            id: The config ID
-            name: Configuration name
-            updated_at: The DB timestamp when this config was last updated at
-            cold_query_result_retention: Retention window for cold query result / state information.
-            cold_storage: Retention window for logs in cold storage. Deleted afterwards.
-            forwarding_configuration: Configuration for forwarding audit logs to external systems
-            hot_storage: Retention window to keep audit logs in hot storage
-            query_result_retention: Length of time to preserve server-side query result caching
-            redact_service_config_id: A redact service config that will be used to redact PII from logs
-            retention: Retention window to store audit logs
-            schema: Audit log field configuration
-            vault_key_id: ID of the Vault key used for signing. If missing, use a default Audit key.
-            vault_service_config_id: A vault service config that will be used to sign logs
-            vault_sign: Enable/disable event signing
-            warm_storage: Retention window for logs in warm storage. Migrated to cold or deleted afterwards.
-        """
-
-        response = self.request.post(
-            "v1beta/config/update",
-            PangeaResponseResult,
-            data={
-                "id": config_id,
-                "name": name,
-                "updated_at": updated_at,
-                # Optionals.
-                "cold_query_result_retention": cold_query_result_retention,
-                "cold_storage": cold_storage,
-                "forwarding_configuration": forwarding_configuration,
-                "hot_storage": hot_storage,
-                "query_result_retention": query_result_retention,
-                "redact_service_config_id": redact_service_config_id,
-                "retention": retention,
-                "schema": schema,
-                "vault_key_id": vault_key_id,
-                "vault_service_config_id": vault_service_config_id,
-                "vault_sign": vault_sign,
-                "warm_storage": warm_storage,
-            },
-        )
-        response.result = TypeAdapter(ServiceConfig).validate_python(response.json["result"])
-        return cast(PangeaResponse[ServiceConfig], response)
-
-    def delete_service_config(self, config_id: str) -> PangeaResponse[ServiceConfig]:
-        """
-        Delete a service config.
-
-        OperationId: audit_post_v1beta_config_delete
-
-        Args:
-            id: The config ID
-        """
-
-        response = self.request.post("v1beta/config/delete", PangeaResponseResult, data={"id": config_id})
-        response.result = TypeAdapter(ServiceConfig).validate_python(response.json["result"])
-        return cast(PangeaResponse[ServiceConfig], response)
-
-    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[ServiceConfigListResult]:
-        """
-        List service configs.
-
-        OperationId: audit_post_v1beta_config_list
-
-        Args:
-            last: Reflected value from a previous response to obtain the next page of results.
-            order: Order results asc(ending) or desc(ending).
-            order_by: Which field to order results by.
-            size: Maximum results to include in the response.
-        """
-
-        return self.request.post(
-            "v1beta/config/list",
-            ServiceConfigListResult,
-            data={"filter": filter, "last": last, "order": order, "order_by": order_by, "size": size},
-        )
-
     def update_published_roots(self, result: SearchResultOutput):
         """Fetches series of published root hashes from Arweave