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

pangea/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "6.4.0"
+__version__ = "6.5.0beta1"
 
 from pangea.asyncio.request import PangeaRequestAsync
 from pangea.config import PangeaConfig

pangea/asyncio/request.py

@@ -10,13 +10,13 @@ import asyncio
 import json
 import time
 from collections.abc import Iterable, Mapping
-from typing import Dict, List, Optional, Sequence, Tuple, Type, Union, cast
+from typing import Dict, List, Optional, Sequence, Tuple, Type, Union, cast, overload
 
 import aiohttp
 from aiohttp import FormData
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 from pydantic_core import to_jsonable_python
-from typing_extensions import Any, TypeAlias, TypeVar, override
+from typing_extensions import Any, Literal, TypeAlias, TypeVar, override
 
 import pangea.exceptions as pe
 from pangea.request import MultipartResponse, PangeaRequestBase
@@ -50,6 +50,24 @@ 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,
@@ -58,18 +76,60 @@ class PangeaRequestAsync(PangeaRequestBase):
         files: Optional[List[Tuple]] = None,
         poll_result: bool = True,
         url: Optional[str] = None,
+        *,
+        pangea_response: Literal[True] = True,
     ) -> PangeaResponse[TResult]:
-        """Makes the POST call to a Pangea Service endpoint.
+        """
+        Makes a POST call to a Pangea Service endpoint.
 
         Args:
-            endpoint(str): The Pangea Service API endpoint.
-            data(dict): The POST body payload object
+            endpoint: The Pangea Service API endpoint.
+            data: 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)
 
@@ -109,9 +169,13 @@ 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: PangeaResponse = PangeaResponse(
+            pangea_response_obj: PangeaResponse = PangeaResponse(
                 requests_response,
                 result_class=result_class,
                 json=multipart_response.pangea_json,
@@ -124,49 +188,110 @@ class PangeaRequestAsync(PangeaRequestBase):
                     json.dumps({"service": self.service, "action": "post", "url": url, "response": json_resp})
                 )
 
-                pangea_response = PangeaResponse(requests_response, result_class=result_class, json=json_resp)
+                pangea_response_obj = 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 = await self._handle_queued_result(pangea_response)
+            pangea_response_obj = await self._handle_queued_result(pangea_response_obj)
 
-        return self._check_response(pangea_response)
+        return self._check_response(pangea_response_obj)
 
-    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.
+    @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.
 
         Args:
-            endpoint(str): The Pangea Service API endpoint.
-            path(str): Additional URL path
+            path: Additional URL path
+            params: Dictionary of querystring data to attach to the request
 
         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, headers=self._headers()) as requests_response:
+        async with self.session.get(url, params=params, headers=self._headers()) as requests_response:
             await self._check_http_errors(requests_response)
-            pangea_response = PangeaResponse(
+
+            if not pangea_response:
+                type_adapter = TypeAdapter(result_class)
+                return type_adapter.validate_python(await requests_response.json())
+
+            pangea_response_obj = 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.json},
+                {"service": self.service, "action": "get", "url": url, "response": pangea_response_obj.json},
                 default=default_encoder,
             )
         )
 
         if check_response is False:
-            return pangea_response
+            return pangea_response_obj
 
-        return self._check_response(pangea_response)
+        return self._check_response(pangea_response_obj)
 
     async def _check_http_errors(self, resp: aiohttp.ClientResponse):
         if resp.status == 503:
@@ -300,6 +425,14 @@ 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,

pangea/asyncio/services/__init__.py

@@ -7,6 +7,7 @@ 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

pangea/asyncio/services/ai_guard.py

@@ -1,12 +1,29 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
-from typing import overload
+from collections.abc import Mapping, Sequence
+from typing import Any, overload
+
+from typing_extensions import Literal, TypeVar
 
 from pangea.asyncio.services.base import ServiceBaseAsync
 from pangea.config import PangeaConfig
 from pangea.response import PangeaResponse
-from pangea.services.ai_guard import LogFields, Message, Overrides, TextGuardResult
+from pangea.services.ai_guard import (
+    AuditDataActivityConfig,
+    ConnectionsConfig,
+    ExtraInfo,
+    GuardResult,
+    LogFields,
+    Message,
+    Overrides,
+    RecipeConfig,
+    ServiceConfig,
+    ServiceConfigFilter,
+    ServiceConfigsPage,
+    TextGuardResult,
+)
+
+_T = TypeVar("_T")
 
 
 class AIGuardAsync(ServiceBaseAsync):
@@ -169,3 +186,157 @@ class AIGuardAsync(ServiceBaseAsync):
                 "log_fields": log_fields,
             },
         )
+
+    async 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
+
+        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
+
+        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 await 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,
+            },
+        )
+
+    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,
+        )