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

pangea/__init__.py

@@ -1,15 +1,7 @@
-__version__ = "6.5.0beta1"
+__version__ = "6.7.0"
 
-from pangea.asyncio.request import PangeaRequestAsync
 from pangea.config import PangeaConfig
 from pangea.file_uploader import FileUploader
-from pangea.request import PangeaRequest
-from pangea.response import PangeaResponse
+from pangea.response import PangeaResponse, PangeaResponseResult, TransferMethod
 
-__all__ = (
-    "FileUploader",
-    "PangeaConfig",
-    "PangeaRequest",
-    "PangeaRequestAsync",
-    "PangeaResponse",
-)
+__all__ = ("FileUploader", "PangeaConfig", "PangeaResponse", "PangeaResponseResult", "TransferMethod")

pangea/_constants.py

@@ -0,0 +1,4 @@
+from __future__ import annotations
+
+MAX_RETRY_DELAY = 8.0
+RETRYABLE_HTTP_CODES = frozenset({500, 502, 503, 504})

pangea/_typing.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from collections.abc import Iterator, Sequence
+
+from typing_extensions import Any, Protocol, SupportsIndex, TypeVar, overload
+
+T = TypeVar("T")
+T_co = TypeVar("T_co", covariant=True)
+
+
+class SequenceNotStr(Protocol[T_co]):
+    """Sequence-like object that isn't str or bytes."""
+
+    @overload
+    def __getitem__(self, index: SupportsIndex, /) -> T_co: ...
+
+    @overload
+    def __getitem__(self, index: slice, /) -> Sequence[T_co]: ...
+
+    def __contains__(self, value: object, /) -> bool: ...
+
+    def __len__(self) -> int: ...
+
+    def __iter__(self) -> Iterator[T_co]: ...
+
+    def index(self, value: Any, start: int = ..., stop: int = ..., /) -> int: ...
+
+    def count(self, value: Any, /) -> int: ...
+
+    def __reversed__(self) -> Iterator[T_co]: ...

pangea/asyncio/__init__.py

@@ -1,2 +1,3 @@
-# ruff: noqa: F401
 from .file_uploader import FileUploaderAsync
+
+__all__ = ("FileUploaderAsync",)

pangea/asyncio/file_uploader.py

@@ -6,9 +6,10 @@ from __future__ import annotations
 import io
 import logging
 
+from pangea import PangeaConfig, TransferMethod
 from pangea.asyncio.request import PangeaRequestAsync
-from pangea.request import PangeaConfig
-from pangea.response import TransferMethod
+
+__all__ = ("FileUploaderAsync",)
 
 
 class FileUploaderAsync:

pangea/asyncio/request.py

@@ -10,19 +10,23 @@ import asyncio
 import json
 import time
 from collections.abc import Iterable, Mapping
-from typing import Dict, List, Optional, Sequence, Tuple, Type, Union, cast, overload
+from random import random
+from typing import Dict, List, Optional, Sequence, Tuple, Type, Union, cast
 
 import aiohttp
 from aiohttp import FormData
-from pydantic import BaseModel, TypeAdapter
+from pydantic import BaseModel
 from pydantic_core import to_jsonable_python
-from typing_extensions import Any, Literal, TypeAlias, TypeVar, override
+from typing_extensions import Any, TypeAlias, TypeVar, override
 
 import pangea.exceptions as pe
+from pangea._constants import MAX_RETRY_DELAY, RETRYABLE_HTTP_CODES
 from pangea.request import MultipartResponse, PangeaRequestBase
 from pangea.response import AttachedFile, PangeaResponse, PangeaResponseResult, ResponseStatus, TransferMethod
 from pangea.utils import default_encoder
 
+__all__ = ("PangeaRequestAsync",)
+
 _FileName: TypeAlias = Union[str, None]
 _FileContent: TypeAlias = Union[str, bytes]
 _FileContentType: TypeAlias = str
@@ -50,24 +54,6 @@ class PangeaRequestAsync(PangeaRequestBase):
     be set in PangeaConfig.
     """
 
-    async def delete(self, endpoint: str) -> None:
-        """
-        Makes a DELETE call to a Pangea endpoint.
-
-        Args:
-            endpoint: The Pangea API endpoint.
-        """
-
-        url = self._url(endpoint)
-
-        self.logger.debug(
-            json.dumps({"service": self.service, "action": "delete", "url": url}, default=default_encoder)
-        )
-
-        requests_response = await self._http_delete(url, headers=self._headers())
-        await self._check_http_errors(requests_response)
-
-    @overload
     async def post(
         self,
         endpoint: str,
@@ -76,60 +62,18 @@ class PangeaRequestAsync(PangeaRequestBase):
         files: Optional[List[Tuple]] = None,
         poll_result: bool = True,
         url: Optional[str] = None,
-        *,
-        pangea_response: Literal[True] = True,
     ) -> PangeaResponse[TResult]:
-        """
-        Makes a POST call to a Pangea Service endpoint.
+        """Makes the POST call to a Pangea Service endpoint.
 
         Args:
-            endpoint: The Pangea Service API endpoint.
-            data: The POST body payload object
+            endpoint(str): The Pangea Service API endpoint.
+            data(dict): The POST body payload object
 
         Returns:
             PangeaResponse which contains the response in its entirety and
                various properties to retrieve individual fields
         """
 
-    @overload
-    async def post(
-        self,
-        endpoint: str,
-        result_class: Type[TResult],
-        data: str | BaseModel | Mapping[str, Any] | None = None,
-        files: Optional[List[Tuple]] = None,
-        poll_result: bool = True,
-        url: Optional[str] = None,
-        *,
-        pangea_response: Literal[False],
-    ) -> TResult:
-        """
-        Makes a POST call to a Pangea Service endpoint.
-
-        Args:
-            endpoint: The Pangea Service API endpoint.
-            data: The POST body payload object
-        """
-
-    async def post(
-        self,
-        endpoint: str,
-        result_class: Type[TResult],
-        data: str | BaseModel | Mapping[str, Any] | None = None,
-        files: Optional[List[Tuple]] = None,
-        poll_result: bool = True,
-        url: Optional[str] = None,
-        *,
-        pangea_response: bool = True,
-    ) -> PangeaResponse[TResult] | TResult:
-        """
-        Makes a POST call to a Pangea Service endpoint.
-
-        Args:
-            endpoint: The Pangea Service API endpoint.
-            data: The POST body payload object
-        """
-
         if isinstance(data, BaseModel):
             data = data.model_dump(exclude_none=True)
 
@@ -169,13 +113,9 @@ class PangeaRequestAsync(PangeaRequestBase):
 
         await self._check_http_errors(requests_response)
 
-        if not pangea_response:
-            type_adapter = TypeAdapter(result_class)
-            return type_adapter.validate_python(await requests_response.json())
-
         if "multipart/form-data" in requests_response.headers.get("content-type", ""):
             multipart_response = await self._process_multipart_response(requests_response)
-            pangea_response_obj: PangeaResponse = PangeaResponse(
+            pangea_response: PangeaResponse = PangeaResponse(
                 requests_response,
                 result_class=result_class,
                 json=multipart_response.pangea_json,
@@ -188,110 +128,49 @@ class PangeaRequestAsync(PangeaRequestBase):
                     json.dumps({"service": self.service, "action": "post", "url": url, "response": json_resp})
                 )
 
-                pangea_response_obj = PangeaResponse(requests_response, result_class=result_class, json=json_resp)
+                pangea_response = PangeaResponse(requests_response, result_class=result_class, json=json_resp)
             except aiohttp.ContentTypeError as e:
                 raise pe.PangeaException(
                     f"Failed to decode json response. {e}. Body: {await requests_response.text()}"
                 ) from e
 
         if poll_result:
-            pangea_response_obj = await self._handle_queued_result(pangea_response_obj)
+            pangea_response = await self._handle_queued_result(pangea_response)
 
-        return self._check_response(pangea_response_obj)
+        return self._check_response(pangea_response)
 
-    @overload
-    async def get(
-        self,
-        path: str,
-        result_class: Type[TResult],
-        check_response: bool = True,
-        *,
-        params: (
-            Mapping[str | bytes | int | float, str | bytes | int | float | Iterable[str | bytes | int | float] | None]
-            | None
-        ) = None,
-        pangea_response: Literal[True] = True,
-    ) -> PangeaResponse[TResult]:
-        """
-        Makes the GET call to a Pangea Service endpoint.
+    async def get(self, path: str, result_class: Type[TResult], check_response: bool = True) -> PangeaResponse[TResult]:
+        """Makes the GET call to a Pangea Service endpoint.
 
         Args:
-            path: Additional URL path
-            params: Dictionary of querystring data to attach to the request
+            endpoint(str): The Pangea Service API endpoint.
+            path(str): Additional URL path
 
         Returns:
             PangeaResponse which contains the response in its entirety and
                various properties to retrieve individual fields
         """
 
-    @overload
-    async def get(
-        self,
-        path: str,
-        result_class: Type[TResult],
-        check_response: bool = True,
-        *,
-        params: (
-            Mapping[str | bytes | int | float, str | bytes | int | float | Iterable[str | bytes | int | float] | None]
-            | None
-        ) = None,
-        pangea_response: Literal[False] = False,
-    ) -> TResult:
-        """
-        Makes the GET call to a Pangea Service endpoint.
-
-        Args:
-            path: Additional URL path
-            params: Dictionary of querystring data to attach to the request
-        """
-
-    async def get(
-        self,
-        path: str,
-        result_class: Type[TResult],
-        check_response: bool = True,
-        *,
-        params: (
-            Mapping[str | bytes | int | float, str | bytes | int | float | Iterable[str | bytes | int | float] | None]
-            | None
-        ) = None,
-        pangea_response: bool = True,
-    ) -> PangeaResponse[TResult] | TResult:
-        """
-        Makes the GET call to a Pangea Service endpoint.
-
-        Args:
-            path: Additional URL path
-            params: Dictionary of querystring data to attach to the request
-            pangea_response: Whether or not the response body follows Pangea's
-              standard response schema
-        """
-
         url = self._url(path)
         self.logger.debug(json.dumps({"service": self.service, "action": "get", "url": url}))
 
-        async with self.session.get(url, params=params, headers=self._headers()) as requests_response:
+        async with self.session.get(url, headers=self._headers()) as requests_response:
             await self._check_http_errors(requests_response)
-
-            if not pangea_response:
-                type_adapter = TypeAdapter(result_class)
-                return type_adapter.validate_python(await requests_response.json())
-
-            pangea_response_obj = PangeaResponse(
+            pangea_response = PangeaResponse(
                 requests_response, result_class=result_class, json=await requests_response.json()
             )
 
         self.logger.debug(
             json.dumps(
-                {"service": self.service, "action": "get", "url": url, "response": pangea_response_obj.json},
+                {"service": self.service, "action": "get", "url": url, "response": pangea_response.json},
                 default=default_encoder,
             )
         )
 
         if check_response is False:
-            return pangea_response_obj
+            return pangea_response
 
-        return self._check_response(pangea_response_obj)
+        return self._check_response(pangea_response)
 
     async def _check_http_errors(self, resp: aiohttp.ClientResponse):
         if resp.status == 503:
@@ -425,14 +304,6 @@ class PangeaRequestAsync(PangeaRequestBase):
         attached_files = await self._get_attached_files(multipart_reader)
         return MultipartResponse(pangea_json, attached_files)  # type: ignore[arg-type]
 
-    async def _http_delete(
-        self,
-        url: str,
-        *,
-        headers: Mapping[str, str | bytes | None] = {},
-    ) -> aiohttp.ClientResponse:
-        return await self.session.delete(url, headers=headers)
-
     async def _http_post(
         self,
         url: str,
@@ -595,13 +466,46 @@ class PangeaRequestAsync(PangeaRequestBase):
 
     @override
     def _init_session(self) -> aiohttp.ClientSession:
-        # retry_config = Retry(
-        #     total=self.config.request_retries,
-        #     backoff_factor=self.config.request_backoff,
-        #     status_forcelist=[500, 502, 503, 504],
-        # )
-        # adapter = HTTPAdapter(max_retries=retry_config)
-        # TODO: Add retry config
-
-        session = aiohttp.ClientSession()
-        return session
+        return aiohttp.ClientSession(middlewares=[self._retry_middleware])
+
+    def _calculate_retry_timeout(self, remaining_retries: int) -> float:
+        max_retries = self.config.request_retries
+        nb_retries = min(max_retries - remaining_retries, 1000)
+        sleep_seconds = min(self.config.request_backoff * pow(2.0, nb_retries), MAX_RETRY_DELAY)
+        jitter = 1 - 0.25 * random()
+        timeout = sleep_seconds * jitter
+        return max(timeout, 0)
+
+    async def _retry_middleware(
+        self, request: aiohttp.ClientRequest, handler: aiohttp.ClientHandlerType
+    ) -> aiohttp.ClientResponse:
+        max_retries = self.config.request_retries
+        request_ids = set[str]()
+        retries_taken = 0
+        for retries_taken in range(max_retries + 1):
+            remaining_retries = max_retries - retries_taken
+
+            if len(request_ids) > 0:
+                request.headers["X-Pangea-Retried-Request-Ids"] = ",".join(request_ids)
+
+            response = await handler(request)
+
+            request_id = response.headers.get("x-request-id")
+            if request_id:
+                request_ids.add(request_id)
+
+            if not response.ok and remaining_retries > 0 and self._should_retry(response):
+                await self._sleep_for_retry(retries_taken=retries_taken, max_retries=max_retries)
+                continue
+
+            break
+
+        return response
+
+    def _should_retry(self, response: aiohttp.ClientResponse) -> bool:
+        return response.status in RETRYABLE_HTTP_CODES
+
+    async def _sleep_for_retry(self, *, retries_taken: int, max_retries: int) -> None:
+        remaining_retries = max_retries - retries_taken
+        timeout = self._calculate_retry_timeout(remaining_retries)
+        await asyncio.sleep(timeout)

pangea/asyncio/services/__init__.py

@@ -1,5 +1,3 @@
-# ruff: noqa: F401
-
 from .ai_guard import AIGuardAsync
 from .audit import AuditAsync
 from .authn import AuthNAsync
@@ -7,9 +5,27 @@ from .authz import AuthZAsync
 from .embargo import EmbargoAsync
 from .file_scan import FileScanAsync
 from .intel import DomainIntelAsync, FileIntelAsync, IpIntelAsync, UrlIntelAsync, UserIntelAsync
-from .management import ManagementAsync
 from .prompt_guard import PromptGuardAsync
 from .redact import RedactAsync
 from .sanitize import SanitizeAsync
 from .share import ShareAsync
 from .vault import VaultAsync
+
+__all__ = (
+    "AIGuardAsync",
+    "AuditAsync",
+    "AuthNAsync",
+    "AuthZAsync",
+    "DomainIntelAsync",
+    "EmbargoAsync",
+    "FileIntelAsync",
+    "FileScanAsync",
+    "IpIntelAsync",
+    "PromptGuardAsync",
+    "RedactAsync",
+    "SanitizeAsync",
+    "ShareAsync",
+    "UrlIntelAsync",
+    "UserIntelAsync",
+    "VaultAsync",
+)

pangea/asyncio/services/ai_guard.py

@@ -1,30 +1,25 @@
 from __future__ import annotations
 
 from collections.abc import Mapping, Sequence
-from typing import Any, overload
+from typing import overload
 
-from typing_extensions import Literal, TypeVar
+from typing_extensions import Any, Literal
 
 from pangea.asyncio.services.base import ServiceBaseAsync
 from pangea.config import PangeaConfig
 from pangea.response import PangeaResponse
 from pangea.services.ai_guard import (
-    AuditDataActivityConfig,
-    ConnectionsConfig,
     ExtraInfo,
     GuardResult,
     LogFields,
+    McpToolsMessage,
     Message,
     Overrides,
-    RecipeConfig,
-    ServiceConfig,
-    ServiceConfigFilter,
-    ServiceConfigsPage,
     TextGuardResult,
+    get_relevant_content,
+    patch_messages,
 )
 
-_T = TypeVar("_T")
-
 
 class AIGuardAsync(ServiceBaseAsync):
     """AI Guard service client.
@@ -100,11 +95,12 @@ class AIGuardAsync(ServiceBaseAsync):
     async 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
@@ -127,6 +123,8 @@ class AIGuardAsync(ServiceBaseAsync):
             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 = await ai_guard.guard_text(messages=[Message(role="user", content="hello world")])
@@ -136,11 +134,12 @@ class AIGuardAsync(ServiceBaseAsync):
         self,
         text: str | None = None,
         *,
-        messages: Sequence[Message] | None = None,
+        messages: Sequence[Message | McpToolsMessage] | None = None,
         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
@@ -166,6 +165,8 @@ class AIGuardAsync(ServiceBaseAsync):
             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 = await ai_guard.guard_text("text")
@@ -174,7 +175,11 @@ class AIGuardAsync(ServiceBaseAsync):
         if text is not None and messages is not None:
             raise ValueError("Exactly one of `text` or `messages` must be given")
 
-        return await self.request.post(
+        if only_relevant_content and messages is not None:
+            original_messages = messages
+            messages, original_indices = get_relevant_content(messages)
+
+        response = await self.request.post(
             "v1/text/guard",
             TextGuardResult,
             data={
@@ -187,6 +192,13 @@ class AIGuardAsync(ServiceBaseAsync):
             },
         )
 
+        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
+
     async def guard(
         self,
         input: Mapping[str, Any],
@@ -204,8 +216,8 @@ class AIGuardAsync(ServiceBaseAsync):
         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]:
@@ -215,7 +227,7 @@ class AIGuardAsync(ServiceBaseAsync):
         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
@@ -234,12 +246,12 @@ class AIGuardAsync(ServiceBaseAsync):
             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 await self.request.post(
-            "v1beta/guard",
+            "v1/guard",
             GuardResult,
             data={
                 "input": input,
@@ -257,86 +269,8 @@ class AIGuardAsync(ServiceBaseAsync):
                 "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,
             },
         )
-
-    async def get_service_config(self, id: str) -> PangeaResponse[ServiceConfig]:
-        """
-        OperationId: ai_guard_post_v1beta_config
-        """
-        return await self.request.post("v1beta/config", data={"id": id}, result_class=ServiceConfig)
-
-    async 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 await self.request.post(
-            "v1beta/config/create",
-            data={
-                "name": name,
-                "id": id,
-                "audit_data_activity": audit_data_activity,
-                "connections": connections,
-                "recipes": recipes,
-            },
-            result_class=ServiceConfig,
-        )
-
-    async 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 await self.request.post(
-            "v1beta/config/update",
-            data={
-                "id": id,
-                "name": name,
-                "audit_data_activity": audit_data_activity,
-                "connections": connections,
-                "recipes": recipes,
-            },
-            result_class=ServiceConfig,
-        )
-
-    async def delete_service_config(self, id: str) -> PangeaResponse[ServiceConfig]:
-        """
-        OperationId: ai_guard_post_v1beta_config_delete
-        """
-        return await self.request.post("v1beta/config/delete", data={"id": id}, result_class=ServiceConfig)
-
-    async 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 await self.request.post(
-            "v1beta/config/list",
-            data={"filter": filter, "last": last, "order": order, "order_by": order_by, "size": size},
-            result_class=ServiceConfigsPage,
-        )