mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Add Purview Middleware (#1142)
* [Py Purview] Purview Python Initial Commit * [Py Purview] Purview Python Minor Fixes * [Py Purview] Purview Python Comment Fixesish * [Py Purview] Purview Python Agent Middleware Done * [Py Purview] Purview Python Agent Middleware Done * [Py Purview] Purview Python Lint Errors * [Py Purview] Purview Python Final Hopefully * [Py Purview] Purview Python Final Hopefully * [Py Purview] Purview Python Fix ReadMe * [Py Purview] Purview Python Fix MyPy * [Py Purview] Purview Python Minor Updates on comments * [Py Purview] Purview Python Fix Build Error --------- Co-authored-by: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com>
This commit is contained in:
committed by
GitHub
Unverified
parent
76ae0a62ac
commit
59da578902
@@ -0,0 +1,22 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from ._exceptions import (
|
||||
PurviewAuthenticationError,
|
||||
PurviewRateLimitError,
|
||||
PurviewRequestError,
|
||||
PurviewServiceError,
|
||||
)
|
||||
from ._middleware import PurviewChatPolicyMiddleware, PurviewPolicyMiddleware
|
||||
from ._settings import PurviewAppLocation, PurviewLocationType, PurviewSettings
|
||||
|
||||
__all__ = [
|
||||
"PurviewAppLocation",
|
||||
"PurviewAuthenticationError",
|
||||
"PurviewChatPolicyMiddleware",
|
||||
"PurviewLocationType",
|
||||
"PurviewPolicyMiddleware",
|
||||
"PurviewRateLimitError",
|
||||
"PurviewRequestError",
|
||||
"PurviewServiceError",
|
||||
"PurviewSettings",
|
||||
]
|
||||
@@ -0,0 +1,126 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import inspect
|
||||
import json
|
||||
from typing import Any, cast
|
||||
|
||||
import httpx
|
||||
from agent_framework import AGENT_FRAMEWORK_USER_AGENT
|
||||
from agent_framework.observability import get_tracer
|
||||
from azure.core.credentials import TokenCredential
|
||||
from azure.core.credentials_async import AsyncTokenCredential
|
||||
|
||||
from ._exceptions import (
|
||||
PurviewAuthenticationError,
|
||||
PurviewRateLimitError,
|
||||
PurviewRequestError,
|
||||
PurviewServiceError,
|
||||
)
|
||||
from ._models import (
|
||||
ContentActivitiesRequest,
|
||||
ContentActivitiesResponse,
|
||||
ProcessContentRequest,
|
||||
ProcessContentResponse,
|
||||
ProtectionScopesRequest,
|
||||
ProtectionScopesResponse,
|
||||
)
|
||||
from ._settings import PurviewSettings
|
||||
|
||||
|
||||
class PurviewClient:
|
||||
"""Async client for calling Graph Purview endpoints.
|
||||
|
||||
Supports both synchronous TokenCredential and asynchronous AsyncTokenCredential implementations.
|
||||
A sync credential will be invoked in a thread to avoid blocking the event loop.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
credential: TokenCredential | AsyncTokenCredential,
|
||||
settings: PurviewSettings,
|
||||
*,
|
||||
timeout: float | None = 10.0,
|
||||
):
|
||||
self._credential: TokenCredential | AsyncTokenCredential = credential
|
||||
self._settings = settings
|
||||
self._graph_uri = settings.graph_base_uri.rstrip("/")
|
||||
self._timeout = timeout
|
||||
self._client = httpx.AsyncClient(timeout=timeout)
|
||||
|
||||
async def close(self) -> None:
|
||||
await self._client.aclose()
|
||||
|
||||
async def _get_token(self, *, tenant_id: str | None = None) -> str:
|
||||
"""Acquire an access token using either async or sync credential."""
|
||||
scopes = self._settings.get_scopes()
|
||||
cred = self._credential
|
||||
token = cred.get_token(*scopes, tenant_id=tenant_id)
|
||||
token = await token if inspect.isawaitable(token) else token
|
||||
return token.token
|
||||
|
||||
@staticmethod
|
||||
def _extract_token_info(token: str) -> dict[str, Any]:
|
||||
parts = token.split(".")
|
||||
if len(parts) < 2:
|
||||
raise ValueError("Invalid JWT token format")
|
||||
payload = parts[1]
|
||||
rem = len(payload) % 4
|
||||
if rem:
|
||||
payload += "=" * (4 - rem)
|
||||
decoded = base64.urlsafe_b64decode(payload)
|
||||
data = json.loads(decoded.decode("utf-8"))
|
||||
return {
|
||||
"user_id": data.get("oid") if data.get("idtyp") == "user" else None,
|
||||
"tenant_id": data.get("tid"),
|
||||
"client_id": data.get("appid"),
|
||||
}
|
||||
|
||||
async def get_user_info_from_token(self, *, tenant_id: str | None = None) -> dict[str, Any]:
|
||||
token = await self._get_token(tenant_id=tenant_id)
|
||||
return self._extract_token_info(token)
|
||||
|
||||
async def process_content(self, request: ProcessContentRequest) -> ProcessContentResponse:
|
||||
with get_tracer().start_as_current_span("purview.process_content"):
|
||||
token = await self._get_token(tenant_id=request.tenant_id)
|
||||
url = f"{self._graph_uri}/users/{request.user_id}/dataSecurityAndGovernance/processContent"
|
||||
return cast(ProcessContentResponse, await self._post(url, request, ProcessContentResponse, token))
|
||||
|
||||
async def get_protection_scopes(self, request: ProtectionScopesRequest) -> ProtectionScopesResponse:
|
||||
with get_tracer().start_as_current_span("purview.get_protection_scopes"):
|
||||
token = await self._get_token()
|
||||
url = f"{self._graph_uri}/users/{request.user_id}/dataSecurityAndGovernance/protectionScopes/compute"
|
||||
return cast(ProtectionScopesResponse, await self._post(url, request, ProtectionScopesResponse, token))
|
||||
|
||||
async def send_content_activities(self, request: ContentActivitiesRequest) -> ContentActivitiesResponse:
|
||||
with get_tracer().start_as_current_span("purview.send_content_activities"):
|
||||
token = await self._get_token()
|
||||
url = f"{self._graph_uri}/users/{request.user_id}/dataSecurityAndGovernance/activities/contentActivities"
|
||||
return cast(ContentActivitiesResponse, await self._post(url, request, ContentActivitiesResponse, token))
|
||||
|
||||
async def _post(self, url: str, model: Any, response_type: type[Any], token: str) -> Any:
|
||||
payload = model.model_dump(by_alias=True, exclude_none=True, mode="json")
|
||||
headers = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"User-Agent": AGENT_FRAMEWORK_USER_AGENT,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
resp = await self._client.post(url, json=payload, headers=headers)
|
||||
if resp.status_code in (401, 403):
|
||||
raise PurviewAuthenticationError(f"Auth failure {resp.status_code}: {resp.text}")
|
||||
if resp.status_code == 429:
|
||||
raise PurviewRateLimitError(f"Rate limited {resp.status_code}: {resp.text}")
|
||||
if resp.status_code not in (200, 201, 202):
|
||||
raise PurviewRequestError(f"Purview request failed {resp.status_code}: {resp.text}")
|
||||
try:
|
||||
data = resp.json()
|
||||
except ValueError:
|
||||
data = {}
|
||||
try:
|
||||
# Prefer pydantic-style model_validate if present, else fall back to constructor.
|
||||
if hasattr(response_type, "model_validate"):
|
||||
return response_type.model_validate(data) # type: ignore[no-any-return]
|
||||
return response_type(**data) # type: ignore[call-arg, no-any-return]
|
||||
except Exception as ex: # pragma: no cover
|
||||
raise PurviewServiceError(f"Failed to deserialize Purview response: {ex}") from ex
|
||||
@@ -0,0 +1,29 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
"""Purview specific exceptions (minimal error shaping)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from agent_framework.exceptions import ServiceResponseException
|
||||
|
||||
__all__ = [
|
||||
"PurviewAuthenticationError",
|
||||
"PurviewRateLimitError",
|
||||
"PurviewRequestError",
|
||||
"PurviewServiceError",
|
||||
]
|
||||
|
||||
|
||||
class PurviewServiceError(ServiceResponseException):
|
||||
"""Base exception for Purview errors."""
|
||||
|
||||
|
||||
class PurviewAuthenticationError(PurviewServiceError):
|
||||
"""Authentication / authorization failure (401/403)."""
|
||||
|
||||
|
||||
class PurviewRateLimitError(PurviewServiceError):
|
||||
"""Rate limiting or throttling (429)."""
|
||||
|
||||
|
||||
class PurviewRequestError(PurviewServiceError):
|
||||
"""Other non-success HTTP errors."""
|
||||
@@ -0,0 +1,168 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
from agent_framework import AgentMiddleware, AgentRunContext, ChatContext, ChatMiddleware
|
||||
from agent_framework._logging import get_logger
|
||||
from azure.core.credentials import TokenCredential
|
||||
from azure.core.credentials_async import AsyncTokenCredential
|
||||
|
||||
from ._client import PurviewClient
|
||||
from ._models import Activity
|
||||
from ._processor import ScopedContentProcessor
|
||||
from ._settings import PurviewSettings
|
||||
|
||||
logger = get_logger("agent_framework.purview")
|
||||
|
||||
|
||||
class PurviewPolicyMiddleware(AgentMiddleware):
|
||||
"""Agent middleware that enforces Purview policies on prompt and response.
|
||||
|
||||
Accepts either a synchronous TokenCredential or an AsyncTokenCredential.
|
||||
|
||||
Usage:
|
||||
|
||||
.. code-block:: python
|
||||
from agent_framework.microsoft import PurviewPolicyMiddleware, PurviewSettings
|
||||
from agent_framework import ChatAgent
|
||||
|
||||
credential = ... # TokenCredential or AsyncTokenCredential
|
||||
settings = PurviewSettings(app_name="My App")
|
||||
agent = ChatAgent(
|
||||
chat_client=client, instructions="...", middleware=[PurviewPolicyMiddleware(credential, settings)]
|
||||
)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
credential: TokenCredential | AsyncTokenCredential,
|
||||
settings: PurviewSettings,
|
||||
) -> None:
|
||||
self._client = PurviewClient(credential, settings)
|
||||
self._processor = ScopedContentProcessor(self._client, settings)
|
||||
self._settings = settings
|
||||
|
||||
async def process(
|
||||
self,
|
||||
context: AgentRunContext,
|
||||
next: Callable[[AgentRunContext], Awaitable[None]],
|
||||
) -> None: # type: ignore[override]
|
||||
resolved_user_id: str | None = None
|
||||
try:
|
||||
# Pre (prompt) check
|
||||
should_block_prompt, resolved_user_id = await self._processor.process_messages(
|
||||
context.messages, Activity.UPLOAD_TEXT
|
||||
)
|
||||
if should_block_prompt:
|
||||
from agent_framework import AgentRunResponse, ChatMessage, Role
|
||||
|
||||
context.result = AgentRunResponse(
|
||||
messages=[ChatMessage(role=Role.SYSTEM, text=self._settings.blocked_prompt_message)]
|
||||
)
|
||||
context.terminate = True
|
||||
return
|
||||
except Exception as ex:
|
||||
# Log and continue if there's an error in the pre-check
|
||||
logger.error(f"Error in Purview policy pre-check: {ex}")
|
||||
|
||||
await next(context)
|
||||
|
||||
try:
|
||||
# Post (response) check only if we have a normal AgentRunResponse
|
||||
# Use the same user_id from the request for the response evaluation
|
||||
if context.result and not context.is_streaming:
|
||||
should_block_response, _ = await self._processor.process_messages(
|
||||
context.result.messages, # type: ignore[union-attr]
|
||||
Activity.UPLOAD_TEXT,
|
||||
user_id=resolved_user_id,
|
||||
)
|
||||
if should_block_response:
|
||||
from agent_framework import AgentRunResponse, ChatMessage, Role
|
||||
|
||||
context.result = AgentRunResponse(
|
||||
messages=[ChatMessage(role=Role.SYSTEM, text=self._settings.blocked_response_message)]
|
||||
)
|
||||
else:
|
||||
# Streaming responses are not supported for post-checks
|
||||
logger.debug("Streaming responses are not supported for Purview policy post-checks")
|
||||
except Exception as ex:
|
||||
# Log and continue if there's an error in the post-check
|
||||
logger.error(f"Error in Purview policy post-check: {ex}")
|
||||
|
||||
|
||||
class PurviewChatPolicyMiddleware(ChatMiddleware):
|
||||
"""Chat middleware variant for Purview policy evaluation.
|
||||
|
||||
This allows users to attach Purview enforcement directly to a chat client
|
||||
|
||||
Behavior:
|
||||
* Pre-chat: evaluates outgoing (user + context) messages as an upload activity
|
||||
and can terminate execution if blocked.
|
||||
* Post-chat: evaluates the received response messages (streaming is not presently supported)
|
||||
and can replace them with a blocked message. Uses the same user_id from the request
|
||||
to ensure consistent user identity throughout the evaluation.
|
||||
|
||||
Usage:
|
||||
|
||||
.. code-block:: python
|
||||
from agent_framework.microsoft import PurviewChatPolicyMiddleware, PurviewSettings
|
||||
from agent_framework import ChatClient
|
||||
|
||||
credential = ... # TokenCredential or AsyncTokenCredential
|
||||
settings = PurviewSettings(app_name="My App")
|
||||
client = ChatClient(..., middleware=[PurviewChatPolicyMiddleware(credential, settings)])
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
credential: TokenCredential | AsyncTokenCredential,
|
||||
settings: PurviewSettings,
|
||||
) -> None:
|
||||
self._client = PurviewClient(credential, settings)
|
||||
self._processor = ScopedContentProcessor(self._client, settings)
|
||||
self._settings = settings
|
||||
|
||||
async def process(
|
||||
self,
|
||||
context: ChatContext,
|
||||
next: Callable[[ChatContext], Awaitable[None]],
|
||||
) -> None: # type: ignore[override]
|
||||
resolved_user_id: str | None = None
|
||||
try:
|
||||
should_block_prompt, resolved_user_id = await self._processor.process_messages(
|
||||
context.messages, Activity.UPLOAD_TEXT
|
||||
)
|
||||
if should_block_prompt:
|
||||
from agent_framework import ChatMessage
|
||||
|
||||
context.result = [ # type: ignore[assignment]
|
||||
ChatMessage(role="system", text=self._settings.blocked_prompt_message)
|
||||
]
|
||||
context.terminate = True
|
||||
return
|
||||
except Exception as ex:
|
||||
logger.error(f"Error in Purview policy pre-check: {ex}")
|
||||
|
||||
await next(context)
|
||||
|
||||
try:
|
||||
# Post (response) evaluation only if non-streaming and we have messages result shape
|
||||
# Use the same user_id from the request for the response evaluation
|
||||
if context.result and not context.is_streaming:
|
||||
result_obj = context.result
|
||||
messages = getattr(result_obj, "messages", None)
|
||||
if messages:
|
||||
should_block_response, _ = await self._processor.process_messages(
|
||||
messages, Activity.UPLOAD_TEXT, user_id=resolved_user_id
|
||||
)
|
||||
if should_block_response:
|
||||
from agent_framework import ChatMessage
|
||||
|
||||
context.result = [ # type: ignore[assignment]
|
||||
ChatMessage(role="system", text=self._settings.blocked_response_message)
|
||||
]
|
||||
else:
|
||||
logger.debug("Streaming responses are not supported for Purview policy post-checks")
|
||||
except Exception as ex:
|
||||
logger.error(f"Error in Purview policy post-check: {ex}")
|
||||
@@ -0,0 +1,992 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Unified Purview model definitions and public export surface."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping, MutableMapping, Sequence
|
||||
from datetime import datetime
|
||||
from enum import Enum, Flag, auto
|
||||
from typing import Any, ClassVar, TypeVar, cast
|
||||
from uuid import uuid4
|
||||
|
||||
from agent_framework._logging import get_logger
|
||||
from agent_framework._serialization import SerializationMixin
|
||||
|
||||
logger = get_logger("agent_framework.purview")
|
||||
|
||||
# --------------------------------------------------------------------------------------
|
||||
# Enums & flag helpers
|
||||
# --------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
class Activity(str, Enum):
|
||||
"""High-level activity types representing user or agent operations."""
|
||||
|
||||
UNKNOWN = "unknown"
|
||||
UPLOAD_TEXT = "uploadText"
|
||||
UPLOAD_FILE = "uploadFile"
|
||||
DOWNLOAD_TEXT = "downloadText"
|
||||
DOWNLOAD_FILE = "downloadFile"
|
||||
|
||||
|
||||
class ProtectionScopeActivities(Flag):
|
||||
"""Flag enumeration of activities used in policy protection scopes."""
|
||||
|
||||
NONE = 0
|
||||
UPLOAD_TEXT = auto()
|
||||
UPLOAD_FILE = auto()
|
||||
DOWNLOAD_TEXT = auto()
|
||||
DOWNLOAD_FILE = auto()
|
||||
UNKNOWN_FUTURE_VALUE = auto()
|
||||
|
||||
def __int__(self) -> int: # pragma: no cover
|
||||
return self.value
|
||||
|
||||
|
||||
FlagT = TypeVar("FlagT", bound=Flag)
|
||||
|
||||
_PROTECTION_SCOPE_ACTIVITIES_MAP: dict[str, ProtectionScopeActivities] = {
|
||||
"none": ProtectionScopeActivities.NONE,
|
||||
"uploadText": ProtectionScopeActivities.UPLOAD_TEXT,
|
||||
"uploadFile": ProtectionScopeActivities.UPLOAD_FILE,
|
||||
"downloadText": ProtectionScopeActivities.DOWNLOAD_TEXT,
|
||||
"downloadFile": ProtectionScopeActivities.DOWNLOAD_FILE,
|
||||
"unknownFutureValue": ProtectionScopeActivities.UNKNOWN_FUTURE_VALUE,
|
||||
}
|
||||
_PROTECTION_SCOPE_ACTIVITIES_SERIALIZE_ORDER: list[tuple[str, ProtectionScopeActivities]] = [
|
||||
("uploadText", ProtectionScopeActivities.UPLOAD_TEXT),
|
||||
("uploadFile", ProtectionScopeActivities.UPLOAD_FILE),
|
||||
("downloadText", ProtectionScopeActivities.DOWNLOAD_TEXT),
|
||||
("downloadFile", ProtectionScopeActivities.DOWNLOAD_FILE),
|
||||
]
|
||||
|
||||
|
||||
def deserialize_flag(
|
||||
value: object, mapping: Mapping[str, FlagT], enum_cls: type[FlagT]
|
||||
) -> FlagT | None: # pragma: no cover
|
||||
"""Deserialize arbitrary input into a flag enum instance."""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, enum_cls):
|
||||
return value
|
||||
if isinstance(value, int):
|
||||
try:
|
||||
return enum_cls(value)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
flag_value = enum_cls(0)
|
||||
parts: list[str] = []
|
||||
|
||||
if isinstance(value, str):
|
||||
raw = value.strip()
|
||||
if not raw:
|
||||
return enum_cls(0)
|
||||
parts.extend([p.strip() for p in raw.split(",") if p.strip()])
|
||||
elif isinstance(value, (list, tuple, set)):
|
||||
for item in value:
|
||||
if isinstance(item, str):
|
||||
parts.extend([p.strip() for p in item.split(",") if p.strip()])
|
||||
elif isinstance(item, enum_cls):
|
||||
flag_value |= item
|
||||
elif isinstance(item, int):
|
||||
try:
|
||||
flag_value |= enum_cls(item)
|
||||
except Exception:
|
||||
logger.warning(f"Failed to convert int {item} to {enum_cls.__name__}")
|
||||
else:
|
||||
return None
|
||||
|
||||
for part in parts:
|
||||
member = mapping.get(part)
|
||||
if member is not None:
|
||||
flag_value |= member
|
||||
|
||||
if flag_value == enum_cls(0):
|
||||
none_member = mapping.get("none")
|
||||
if none_member is not None:
|
||||
return none_member # type: ignore[return-value,index]
|
||||
return flag_value
|
||||
|
||||
|
||||
def serialize_flag(
|
||||
flag_value: Flag | int | None, ordered_parts: Sequence[tuple[str, Flag]]
|
||||
) -> str | None: # pragma: no cover
|
||||
"""Serialize a flag enum (or int) into a stable, comma-separated string."""
|
||||
if flag_value is None:
|
||||
return None
|
||||
if isinstance(flag_value, int):
|
||||
if flag_value == 0:
|
||||
return "none"
|
||||
int_parts: list[str] = []
|
||||
for name, member in ordered_parts:
|
||||
if flag_value & member.value:
|
||||
int_parts.append(name)
|
||||
return ",".join(int_parts) if int_parts else "none"
|
||||
if not isinstance(flag_value, Flag):
|
||||
return None
|
||||
if flag_value.value == 0:
|
||||
return "none"
|
||||
parts: list[str] = []
|
||||
for name, member in ordered_parts:
|
||||
if flag_value & member:
|
||||
parts.append(name)
|
||||
return ",".join(parts) if parts else "none"
|
||||
|
||||
|
||||
class DlpAction(str, Enum):
|
||||
BLOCK_ACCESS = "blockAccess"
|
||||
OTHER = "other"
|
||||
|
||||
|
||||
class RestrictionAction(str, Enum):
|
||||
BLOCK = "block"
|
||||
OTHER = "other"
|
||||
|
||||
|
||||
class ProtectionScopeState(str, Enum):
|
||||
NOT_MODIFIED = "notModified"
|
||||
MODIFIED = "modified"
|
||||
UNKNOWN_FUTURE_VALUE = "unknownFutureValue"
|
||||
|
||||
|
||||
class ExecutionMode(str, Enum):
|
||||
EVALUATE_INLINE = "evaluateInline"
|
||||
EVALUATE_OFFLINE = "evaluateOffline"
|
||||
UNKNOWN_FUTURE_VALUE = "unknownFutureValue"
|
||||
|
||||
|
||||
class PolicyPivotProperty(str, Enum):
|
||||
NONE = "none"
|
||||
ACTIVITY = "activity"
|
||||
LOCATION = "location"
|
||||
UNKNOWN_FUTURE_VALUE = "unknownFutureValue"
|
||||
|
||||
|
||||
def translate_activity(activity: Activity) -> ProtectionScopeActivities:
|
||||
mapping = {
|
||||
Activity.UNKNOWN: ProtectionScopeActivities.NONE,
|
||||
Activity.UPLOAD_TEXT: ProtectionScopeActivities.UPLOAD_TEXT,
|
||||
Activity.UPLOAD_FILE: ProtectionScopeActivities.UPLOAD_FILE,
|
||||
Activity.DOWNLOAD_TEXT: ProtectionScopeActivities.DOWNLOAD_TEXT,
|
||||
Activity.DOWNLOAD_FILE: ProtectionScopeActivities.DOWNLOAD_FILE,
|
||||
}
|
||||
return mapping.get(activity, ProtectionScopeActivities.UNKNOWN_FUTURE_VALUE)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------
|
||||
# Simple value models
|
||||
# --------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _AliasSerializable(SerializationMixin):
|
||||
"""Base class adding alias mapping + pydantic-compat helpers.
|
||||
|
||||
Each subclass can define ``_ALIASES`` mapping internal attribute name -> external serialized key.
|
||||
``to_dict`` will emit external keys; ``from_dict`` (via ``__init__`` preprocessing) accepts either form.
|
||||
|
||||
Provides light-weight compatibility helpers ``model_dump`` / ``model_validate``
|
||||
"""
|
||||
|
||||
_ALIASES: ClassVar[dict[str, str]] = {}
|
||||
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
# Normalize alias keys -> internal names across the entire class hierarchy
|
||||
# Collect all aliases from parent classes too
|
||||
all_aliases: dict[str, str] = {}
|
||||
for cls in type(self).__mro__:
|
||||
if hasattr(cls, "_ALIASES") and isinstance(cls._ALIASES, dict):
|
||||
for internal, external in cls._ALIASES.items():
|
||||
if external not in all_aliases:
|
||||
all_aliases[external] = internal
|
||||
|
||||
# Normalize all aliased keys in kwargs
|
||||
for external, internal in all_aliases.items():
|
||||
if external in kwargs and internal not in kwargs:
|
||||
kwargs[internal] = kwargs.pop(external)
|
||||
|
||||
# Set normalized kwargs as attributes
|
||||
# This will overwrite any None values that child __init__ may have set from default params
|
||||
for k, v in kwargs.items():
|
||||
setattr(self, k, v)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Compatibility helpers
|
||||
# ------------------------------------------------------------------
|
||||
def model_dump(self, *, by_alias: bool = True, exclude_none: bool = True, **_: Any) -> dict[str, Any]:
|
||||
# Use self.to_dict() to get alias translation
|
||||
d = self.to_dict(exclude_none=exclude_none)
|
||||
# If by_alias=False, translate external -> internal (rarely needed; default True)
|
||||
if not by_alias and self._ALIASES:
|
||||
reverse = {v: k for k, v in self._ALIASES.items()}
|
||||
translated: dict[str, Any] = {}
|
||||
for k, v in d.items():
|
||||
translated[reverse.get(k, k)] = v
|
||||
return translated
|
||||
return d
|
||||
|
||||
def model_dump_json(self, *, by_alias: bool = True, exclude_none: bool = True, **kwargs: Any) -> str:
|
||||
import json
|
||||
|
||||
return json.dumps(self.model_dump(by_alias=by_alias, exclude_none=exclude_none, **kwargs))
|
||||
|
||||
@classmethod
|
||||
def model_validate(cls, value: MutableMapping[str, Any]) -> _AliasSerializable: # type: ignore[name-defined]
|
||||
return cls(**value)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Override to handle alias emission
|
||||
# ------------------------------------------------------------------
|
||||
def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: # type: ignore[override]
|
||||
base = SerializationMixin.to_dict(self, exclude=exclude, exclude_none=exclude_none)
|
||||
|
||||
# For Graph API models, remove the auto-generated 'type' field if it's in DEFAULT_EXCLUDE
|
||||
if "type" in self.DEFAULT_EXCLUDE:
|
||||
base.pop("type", None)
|
||||
|
||||
# Collect all aliases from class hierarchy
|
||||
all_aliases: dict[str, str] = {}
|
||||
for cls in type(self).__mro__:
|
||||
if hasattr(cls, "_ALIASES") and isinstance(cls._ALIASES, dict):
|
||||
# Parent aliases first (will be overridden by child if same key)
|
||||
for internal, external in cls._ALIASES.items():
|
||||
if internal not in all_aliases:
|
||||
all_aliases[internal] = external
|
||||
|
||||
if not all_aliases:
|
||||
return base
|
||||
|
||||
# Translate internal -> external keys (except 'type' reserved)
|
||||
translated: dict[str, Any] = {}
|
||||
for k, v in base.items():
|
||||
if k == "type":
|
||||
translated[k] = v
|
||||
continue
|
||||
external = all_aliases.get(k, k)
|
||||
translated[external] = v
|
||||
return translated
|
||||
|
||||
|
||||
class PolicyLocation(_AliasSerializable):
|
||||
_ALIASES: ClassVar[dict[str, str]] = {"data_type": "@odata.type"}
|
||||
DEFAULT_EXCLUDE: ClassVar[set[str]] = {"type"} # Exclude auto-generated type field for Graph API
|
||||
|
||||
def __init__(self, data_type: str | None = None, value: str | None = None, **kwargs: Any) -> None:
|
||||
# Extract aliased values from kwargs
|
||||
if "@odata.type" in kwargs:
|
||||
data_type = kwargs["@odata.type"]
|
||||
|
||||
# Call parent without explicit params with aliases
|
||||
super().__init__(**kwargs)
|
||||
self.data_type = data_type
|
||||
self.value = value
|
||||
|
||||
|
||||
class ActivityMetadata(_AliasSerializable):
|
||||
_ALIASES: ClassVar[dict[str, str]] = {"activity": "activity"}
|
||||
|
||||
def __init__(self, activity: Activity, **kwargs: Any) -> None:
|
||||
super().__init__(activity=activity, **kwargs)
|
||||
self.activity = activity
|
||||
|
||||
|
||||
class OperatingSystemSpecifications(_AliasSerializable):
|
||||
_ALIASES: ClassVar[dict[str, str]] = {
|
||||
"operating_system_platform": "operatingSystemPlatform",
|
||||
"operating_system_version": "operatingSystemVersion",
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
operating_system_platform: str | None = None,
|
||||
operating_system_version: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
# Extract aliased values from kwargs
|
||||
if "operatingSystemPlatform" in kwargs:
|
||||
operating_system_platform = kwargs["operatingSystemPlatform"]
|
||||
if "operatingSystemVersion" in kwargs:
|
||||
operating_system_version = kwargs["operatingSystemVersion"]
|
||||
|
||||
# Call parent without explicit params with aliases
|
||||
super().__init__(**kwargs)
|
||||
self.operating_system_platform = operating_system_platform
|
||||
self.operating_system_version = operating_system_version
|
||||
|
||||
|
||||
class DeviceMetadata(_AliasSerializable):
|
||||
_ALIASES: ClassVar[dict[str, str]] = {
|
||||
"ip_address": "ipAddress",
|
||||
"operating_system_specifications": "operatingSystemSpecifications",
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ip_address: str | None = None,
|
||||
operating_system_specifications: OperatingSystemSpecifications | MutableMapping[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
# Extract aliased values from kwargs
|
||||
if "ipAddress" in kwargs:
|
||||
ip_address = kwargs["ipAddress"]
|
||||
if "operatingSystemSpecifications" in kwargs:
|
||||
operating_system_specifications = kwargs["operatingSystemSpecifications"]
|
||||
|
||||
# Convert nested objects
|
||||
if isinstance(operating_system_specifications, MutableMapping):
|
||||
operating_system_specifications = OperatingSystemSpecifications(**operating_system_specifications)
|
||||
|
||||
# Call parent without explicit params with aliases
|
||||
super().__init__(**kwargs)
|
||||
self.ip_address = ip_address
|
||||
self.operating_system_specifications = operating_system_specifications
|
||||
|
||||
|
||||
class IntegratedAppMetadata(_AliasSerializable):
|
||||
def __init__(self, name: str | None = None, version: str | None = None, **kwargs: Any) -> None:
|
||||
super().__init__(name=name, version=version, **kwargs)
|
||||
self.name = name
|
||||
self.version = version
|
||||
|
||||
|
||||
class ProtectedAppMetadata(_AliasSerializable):
|
||||
_ALIASES: ClassVar[dict[str, str]] = {"application_location": "applicationLocation"}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str | None = None,
|
||||
version: str | None = None,
|
||||
application_location: PolicyLocation | MutableMapping[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
# Extract aliased values from kwargs
|
||||
if "applicationLocation" in kwargs:
|
||||
application_location = kwargs["applicationLocation"]
|
||||
|
||||
# Convert nested objects
|
||||
if isinstance(application_location, MutableMapping):
|
||||
application_location = PolicyLocation(**application_location)
|
||||
|
||||
# Call parent without explicit params with aliases
|
||||
super().__init__(**kwargs)
|
||||
self.name = name
|
||||
self.version = version
|
||||
self.application_location = application_location # type: ignore[assignment]
|
||||
|
||||
|
||||
class DlpActionInfo(_AliasSerializable):
|
||||
_ALIASES: ClassVar[dict[str, str]] = {"restriction_action": "restrictionAction"}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
action: DlpAction | None = None,
|
||||
restriction_action: RestrictionAction | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
# Extract aliased values from kwargs
|
||||
if "restrictionAction" in kwargs:
|
||||
restriction_action = kwargs["restrictionAction"]
|
||||
|
||||
# Call parent without explicit params with aliases
|
||||
super().__init__(**kwargs)
|
||||
self.action = action
|
||||
self.restriction_action = restriction_action
|
||||
|
||||
|
||||
class AccessedResourceDetails(_AliasSerializable):
|
||||
_ALIASES: ClassVar[dict[str, str]] = {
|
||||
"label_id": "labelId",
|
||||
"access_type": "accessType",
|
||||
"is_cross_prompt_injection_detected": "isCrossPromptInjectionDetected",
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
identifier: str | None = None,
|
||||
name: str | None = None,
|
||||
url: str | None = None,
|
||||
label_id: str | None = None,
|
||||
access_type: str | None = None,
|
||||
status: str | None = None,
|
||||
is_cross_prompt_injection_detected: bool | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
# Extract aliased values from kwargs
|
||||
if "labelId" in kwargs:
|
||||
label_id = kwargs["labelId"]
|
||||
if "accessType" in kwargs:
|
||||
access_type = kwargs["accessType"]
|
||||
if "isCrossPromptInjectionDetected" in kwargs:
|
||||
is_cross_prompt_injection_detected = kwargs["isCrossPromptInjectionDetected"]
|
||||
|
||||
# Call parent without explicit params with aliases
|
||||
super().__init__(**kwargs)
|
||||
self.identifier = identifier
|
||||
self.name = name
|
||||
self.url = url
|
||||
self.label_id = label_id
|
||||
self.access_type = access_type
|
||||
self.status = status
|
||||
self.is_cross_prompt_injection_detected = is_cross_prompt_injection_detected
|
||||
|
||||
|
||||
class AiInteractionPlugin(_AliasSerializable):
|
||||
def __init__(
|
||||
self,
|
||||
identifier: str | None = None,
|
||||
name: str | None = None,
|
||||
version: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
super().__init__(identifier=identifier, name=name, version=version, **kwargs)
|
||||
self.identifier = identifier
|
||||
self.name = name
|
||||
self.version = version
|
||||
|
||||
|
||||
class AiAgentInfo(_AliasSerializable):
|
||||
def __init__(
|
||||
self,
|
||||
identifier: str | None = None,
|
||||
name: str | None = None,
|
||||
version: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
super().__init__(identifier=identifier, name=name, version=version, **kwargs)
|
||||
self.identifier = identifier
|
||||
self.name = name
|
||||
self.version = version
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------
|
||||
# Content models
|
||||
# --------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
class GraphDataTypeBase(_AliasSerializable):
|
||||
_ALIASES: ClassVar[dict[str, str]] = {"data_type": "@odata.type"}
|
||||
# Exclude the auto-generated 'type' field - Graph API uses @odata.type instead
|
||||
DEFAULT_EXCLUDE: ClassVar[set[str]] = {"type"}
|
||||
|
||||
def __init__(self, data_type: str, **kwargs: Any) -> None:
|
||||
super().__init__(data_type=data_type, **kwargs)
|
||||
self.data_type = data_type
|
||||
|
||||
|
||||
class ContentBase(GraphDataTypeBase):
|
||||
pass
|
||||
|
||||
|
||||
class PurviewTextContent(ContentBase):
|
||||
def __init__(self, data: str, data_type: str = "microsoft.graph.textContent", **kwargs: Any) -> None:
|
||||
super().__init__(data_type=data_type, **kwargs)
|
||||
self.data = data
|
||||
|
||||
|
||||
class PurviewBinaryContent(ContentBase):
|
||||
def __init__(self, data: bytes, data_type: str = "microsoft.graph.binaryContent", **kwargs: Any) -> None:
|
||||
super().__init__(data_type=data_type, **kwargs)
|
||||
self.data = data
|
||||
|
||||
def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: # type: ignore[override]
|
||||
import base64
|
||||
|
||||
base = super().to_dict(exclude=exclude, exclude_none=exclude_none)
|
||||
# Ensure bytes encoded as base64 string like pydantic
|
||||
data_bytes = getattr(self, "data", b"") or b""
|
||||
base["data"] = base64.b64encode(data_bytes).decode("utf-8")
|
||||
return base
|
||||
|
||||
|
||||
class ProcessConversationMetadata(GraphDataTypeBase):
|
||||
_ALIASES: ClassVar[dict[str, str]] = {
|
||||
"correlation_id": "correlationId",
|
||||
"sequence_number": "sequenceNumber",
|
||||
"is_truncated": "isTruncated",
|
||||
"created_date_time": "createdDateTime",
|
||||
"modified_date_time": "modifiedDateTime",
|
||||
"parent_message_id": "parentMessageId",
|
||||
"accessed_resources": "accessedResources_v2",
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
identifier: str | None = None,
|
||||
content: PurviewTextContent | PurviewBinaryContent | ContentBase | MutableMapping[str, Any] | None = None,
|
||||
name: str | None = None,
|
||||
is_truncated: bool | None = None,
|
||||
data_type: str = "microsoft.graph.processConversationMetadata", # emitted via base
|
||||
correlation_id: str | None = None,
|
||||
sequence_number: int | None = None,
|
||||
length: int | None = None,
|
||||
created_date_time: datetime | None = None,
|
||||
modified_date_time: datetime | None = None,
|
||||
parent_message_id: str | None = None,
|
||||
accessed_resources: list[AccessedResourceDetails | MutableMapping[str, Any]] | None = None,
|
||||
plugins: list[AiInteractionPlugin | MutableMapping[str, Any]] | None = None,
|
||||
agents: list[AiAgentInfo | MutableMapping[str, Any]] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
# Extract aliased values from kwargs
|
||||
if "correlationId" in kwargs:
|
||||
correlation_id = kwargs["correlationId"]
|
||||
if "sequenceNumber" in kwargs:
|
||||
sequence_number = kwargs["sequenceNumber"]
|
||||
if "isTruncated" in kwargs:
|
||||
is_truncated = kwargs["isTruncated"]
|
||||
if "createdDateTime" in kwargs:
|
||||
created_date_time = kwargs["createdDateTime"]
|
||||
if "modifiedDateTime" in kwargs:
|
||||
modified_date_time = kwargs["modifiedDateTime"]
|
||||
if "parentMessageId" in kwargs:
|
||||
parent_message_id = kwargs["parentMessageId"]
|
||||
if "accessedResources_v2" in kwargs:
|
||||
accessed_resources = kwargs["accessedResources_v2"]
|
||||
|
||||
# Convert nested objects
|
||||
if isinstance(content, MutableMapping):
|
||||
# determine by type? fall back to text content
|
||||
c_type = content.get("@odata.type") or content.get("data_type")
|
||||
if c_type and "binary" in str(c_type):
|
||||
content = PurviewBinaryContent(**content) # type: ignore[arg-type]
|
||||
else:
|
||||
content = PurviewTextContent(**content) # type: ignore[arg-type]
|
||||
accessed_list: list[AccessedResourceDetails] | None = None
|
||||
if accessed_resources:
|
||||
accessed_list = [
|
||||
ar if isinstance(ar, AccessedResourceDetails) else AccessedResourceDetails(**ar)
|
||||
for ar in accessed_resources
|
||||
]
|
||||
plugin_list: list[AiInteractionPlugin] | None = None
|
||||
if plugins:
|
||||
plugin_list = [p if isinstance(p, AiInteractionPlugin) else AiInteractionPlugin(**p) for p in plugins]
|
||||
agent_list: list[AiAgentInfo] | None = None
|
||||
if agents:
|
||||
agent_list = [a if isinstance(a, AiAgentInfo) else AiAgentInfo(**a) for a in agents]
|
||||
|
||||
# Call parent without explicit params with aliases
|
||||
super().__init__(data_type=data_type, **kwargs)
|
||||
self.identifier = identifier
|
||||
self.content = content # type: ignore[assignment]
|
||||
self.name = name
|
||||
self.correlation_id = correlation_id
|
||||
self.sequence_number = sequence_number
|
||||
self.length = length
|
||||
self.is_truncated = is_truncated
|
||||
self.created_date_time = created_date_time
|
||||
self.modified_date_time = modified_date_time
|
||||
self.parent_message_id = parent_message_id
|
||||
self.accessed_resources = accessed_list
|
||||
self.plugins = plugin_list
|
||||
self.agents = agent_list
|
||||
|
||||
|
||||
class ContentToProcess(_AliasSerializable):
|
||||
_ALIASES: ClassVar[dict[str, str]] = {
|
||||
"content_entries": "contentEntries",
|
||||
"activity_metadata": "activityMetadata",
|
||||
"device_metadata": "deviceMetadata",
|
||||
"integrated_app_metadata": "integratedAppMetadata",
|
||||
"protected_app_metadata": "protectedAppMetadata",
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
content_entries: list[ProcessConversationMetadata | MutableMapping[str, Any]],
|
||||
activity_metadata: ActivityMetadata | MutableMapping[str, Any],
|
||||
device_metadata: DeviceMetadata | MutableMapping[str, Any],
|
||||
integrated_app_metadata: IntegratedAppMetadata | MutableMapping[str, Any],
|
||||
protected_app_metadata: ProtectedAppMetadata | MutableMapping[str, Any],
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
# Extract aliased values from kwargs
|
||||
if "contentEntries" in kwargs:
|
||||
content_entries = kwargs["contentEntries"]
|
||||
if "activityMetadata" in kwargs:
|
||||
activity_metadata = kwargs["activityMetadata"]
|
||||
if "deviceMetadata" in kwargs:
|
||||
device_metadata = kwargs["deviceMetadata"]
|
||||
if "integratedAppMetadata" in kwargs:
|
||||
integrated_app_metadata = kwargs["integratedAppMetadata"]
|
||||
if "protectedAppMetadata" in kwargs:
|
||||
protected_app_metadata = kwargs["protectedAppMetadata"]
|
||||
|
||||
# Convert nested objects
|
||||
entries = [
|
||||
e if isinstance(e, ProcessConversationMetadata) else ProcessConversationMetadata(**e)
|
||||
for e in content_entries
|
||||
]
|
||||
if isinstance(activity_metadata, MutableMapping):
|
||||
activity_metadata = ActivityMetadata(**activity_metadata)
|
||||
if isinstance(device_metadata, MutableMapping):
|
||||
device_metadata = DeviceMetadata(**device_metadata)
|
||||
if isinstance(integrated_app_metadata, MutableMapping):
|
||||
integrated_app_metadata = IntegratedAppMetadata(**integrated_app_metadata)
|
||||
if isinstance(protected_app_metadata, MutableMapping):
|
||||
protected_app_metadata = ProtectedAppMetadata(**protected_app_metadata)
|
||||
|
||||
# Call parent without explicit params with aliases
|
||||
super().__init__(**kwargs)
|
||||
self.content_entries = entries
|
||||
self.activity_metadata = activity_metadata # type: ignore[assignment]
|
||||
self.device_metadata = device_metadata # type: ignore[assignment]
|
||||
self.integrated_app_metadata = integrated_app_metadata # type: ignore[assignment]
|
||||
self.protected_app_metadata = protected_app_metadata # type: ignore[assignment]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------
|
||||
# Request models
|
||||
# --------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ProcessContentRequest(_AliasSerializable):
|
||||
_ALIASES: ClassVar[dict[str, str]] = {"content_to_process": "contentToProcess"}
|
||||
DEFAULT_EXCLUDE: ClassVar[set[str]] = {"user_id", "tenant_id", "correlation_id", "process_inline"}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
content_to_process: ContentToProcess | MutableMapping[str, Any],
|
||||
user_id: str,
|
||||
tenant_id: str,
|
||||
correlation_id: str | None = None,
|
||||
process_inline: bool | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
# Extract aliased values from kwargs
|
||||
if "contentToProcess" in kwargs:
|
||||
content_to_process = kwargs["contentToProcess"]
|
||||
|
||||
# Convert nested objects
|
||||
if isinstance(content_to_process, MutableMapping):
|
||||
content_to_process = ContentToProcess(**content_to_process)
|
||||
|
||||
# Call parent without explicit params with aliases
|
||||
super().__init__(**kwargs)
|
||||
self.content_to_process = content_to_process # type: ignore[assignment]
|
||||
self.user_id = user_id
|
||||
self.tenant_id = tenant_id
|
||||
self.correlation_id = correlation_id
|
||||
self.process_inline = process_inline
|
||||
|
||||
|
||||
class ProtectionScopesRequest(_AliasSerializable):
|
||||
DEFAULT_EXCLUDE: ClassVar[set[str]] = {"user_id", "tenant_id", "correlation_id", "scope_identifier"}
|
||||
_ALIASES: ClassVar[dict[str, str]] = {
|
||||
"pivot_on": "pivotOn",
|
||||
"device_metadata": "deviceMetadata",
|
||||
"integrated_app_metadata": "integratedAppMetadata",
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
user_id: str,
|
||||
tenant_id: str,
|
||||
activities: ProtectionScopeActivities | str | int | Sequence[str] | None = None,
|
||||
locations: list[PolicyLocation | MutableMapping[str, Any]] | None = None,
|
||||
pivot_on: PolicyPivotProperty | None = None,
|
||||
device_metadata: DeviceMetadata | MutableMapping[str, Any] | None = None,
|
||||
integrated_app_metadata: IntegratedAppMetadata | MutableMapping[str, Any] | None = None,
|
||||
correlation_id: str | None = None,
|
||||
scope_identifier: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
# Extract aliased values from kwargs
|
||||
if "pivotOn" in kwargs:
|
||||
pivot_on = kwargs["pivotOn"]
|
||||
if "deviceMetadata" in kwargs:
|
||||
device_metadata = kwargs["deviceMetadata"]
|
||||
if "integratedAppMetadata" in kwargs:
|
||||
integrated_app_metadata = kwargs["integratedAppMetadata"]
|
||||
|
||||
# Deserialize activities flag
|
||||
if not isinstance(activities, ProtectionScopeActivities) and activities is not None:
|
||||
activities = deserialize_flag(activities, _PROTECTION_SCOPE_ACTIVITIES_MAP, ProtectionScopeActivities)
|
||||
|
||||
# Convert nested objects
|
||||
if locations:
|
||||
locations = [loc if isinstance(loc, PolicyLocation) else PolicyLocation(**loc) for loc in locations]
|
||||
if isinstance(device_metadata, MutableMapping):
|
||||
device_metadata = DeviceMetadata(**device_metadata)
|
||||
if isinstance(integrated_app_metadata, MutableMapping):
|
||||
integrated_app_metadata = IntegratedAppMetadata(**integrated_app_metadata)
|
||||
|
||||
# Call parent without explicit params with aliases
|
||||
super().__init__(**kwargs)
|
||||
self.user_id = user_id
|
||||
self.tenant_id = tenant_id
|
||||
self.activities = activities # type: ignore[assignment]
|
||||
self.locations = locations
|
||||
self.pivot_on = pivot_on
|
||||
self.device_metadata = device_metadata
|
||||
self.integrated_app_metadata = integrated_app_metadata
|
||||
self.correlation_id = correlation_id
|
||||
self.scope_identifier = scope_identifier
|
||||
|
||||
def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: # type: ignore[override]
|
||||
# Get base dict (activities will be missing because Flag isn't JSON-serializable)
|
||||
base = super().to_dict(exclude=exclude, exclude_none=exclude_none)
|
||||
|
||||
# Manually serialize activities flag if present and not excluded
|
||||
if self.activities is not None or not exclude_none:
|
||||
if self.activities is not None:
|
||||
base["activities"] = serialize_flag(self.activities, _PROTECTION_SCOPE_ACTIVITIES_SERIALIZE_ORDER)
|
||||
elif not exclude_none:
|
||||
base["activities"] = None
|
||||
|
||||
return base
|
||||
|
||||
|
||||
class ContentActivitiesRequest(_AliasSerializable):
|
||||
_ALIASES: ClassVar[dict[str, str]] = {
|
||||
"user_id": "userId",
|
||||
"scope_identifier": "scopeIdentifier",
|
||||
"content_to_process": "contentMetadata",
|
||||
}
|
||||
DEFAULT_EXCLUDE: ClassVar[set[str]] = {"tenant_id", "correlation_id"}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
user_id: str,
|
||||
content_to_process: ContentToProcess | MutableMapping[str, Any],
|
||||
tenant_id: str,
|
||||
id: str | None = None,
|
||||
scope_identifier: str | None = None,
|
||||
correlation_id: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
# Extract aliased values from kwargs
|
||||
if "userId" in kwargs:
|
||||
user_id = kwargs["userId"]
|
||||
if "scopeIdentifier" in kwargs:
|
||||
scope_identifier = kwargs["scopeIdentifier"]
|
||||
if "contentMetadata" in kwargs:
|
||||
content_to_process = kwargs["contentMetadata"]
|
||||
|
||||
# Convert nested objects
|
||||
if isinstance(content_to_process, MutableMapping):
|
||||
content_to_process = ContentToProcess(**content_to_process)
|
||||
|
||||
# Call parent without explicit params with aliases
|
||||
super().__init__(**kwargs)
|
||||
self.id = id or str(uuid4())
|
||||
self.user_id = user_id
|
||||
self.content_to_process = content_to_process # type: ignore[assignment]
|
||||
self.tenant_id = tenant_id
|
||||
self.scope_identifier = scope_identifier
|
||||
self.correlation_id = correlation_id
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------
|
||||
# Response models
|
||||
# --------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ErrorDetails(_AliasSerializable):
|
||||
def __init__(self, code: str | None = None, message: str | None = None, **kwargs: Any) -> None:
|
||||
super().__init__(code=code, message=message, **kwargs)
|
||||
self.code = code
|
||||
self.message = message
|
||||
|
||||
|
||||
class ProcessingError(_AliasSerializable):
|
||||
def __init__(self, message: str | None = None, **kwargs: Any) -> None:
|
||||
super().__init__(message=message, **kwargs)
|
||||
self.message = message
|
||||
|
||||
|
||||
class ProcessContentResponse(_AliasSerializable):
|
||||
_ALIASES: ClassVar[dict[str, str]] = {
|
||||
"protection_scope_state": "protectionScopeState",
|
||||
"policy_actions": "policyActions",
|
||||
"processing_errors": "processingErrors",
|
||||
}
|
||||
|
||||
id: str | None
|
||||
protection_scope_state: ProtectionScopeState | None
|
||||
policy_actions: list[DlpActionInfo] | None
|
||||
processing_errors: list[ProcessingError] | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
id: str | None = None,
|
||||
protection_scope_state: ProtectionScopeState | None = None,
|
||||
policy_actions: list[DlpActionInfo | MutableMapping[str, Any]] | None = None,
|
||||
processing_errors: list[ProcessingError | MutableMapping[str, Any]] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
# Extract aliased values from kwargs
|
||||
if "protectionScopeState" in kwargs:
|
||||
protection_scope_state = kwargs["protectionScopeState"]
|
||||
if "policyActions" in kwargs:
|
||||
policy_actions = kwargs["policyActions"]
|
||||
if "processingErrors" in kwargs:
|
||||
processing_errors = kwargs["processingErrors"]
|
||||
|
||||
# Convert to objects
|
||||
converted_policy_actions: list[DlpActionInfo] | None = None
|
||||
if policy_actions is not None:
|
||||
converted_policy_actions = cast(
|
||||
list[DlpActionInfo],
|
||||
[p if isinstance(p, DlpActionInfo) else DlpActionInfo(**p) for p in policy_actions],
|
||||
)
|
||||
|
||||
converted_processing_errors: list[ProcessingError] | None = None
|
||||
if processing_errors is not None:
|
||||
converted_processing_errors = cast(
|
||||
list[ProcessingError],
|
||||
[pe if isinstance(pe, ProcessingError) else ProcessingError(**pe) for pe in processing_errors],
|
||||
)
|
||||
|
||||
# Call parent without explicit params with aliases
|
||||
super().__init__(**kwargs)
|
||||
self.id = id
|
||||
self.protection_scope_state = protection_scope_state
|
||||
self.policy_actions = converted_policy_actions
|
||||
self.processing_errors = converted_processing_errors
|
||||
|
||||
|
||||
class PolicyScope(_AliasSerializable):
|
||||
_ALIASES: ClassVar[dict[str, str]] = {"policy_actions": "policyActions", "execution_mode": "executionMode"}
|
||||
|
||||
activities: ProtectionScopeActivities | None
|
||||
locations: list[PolicyLocation] | None
|
||||
policy_actions: list[DlpActionInfo] | None
|
||||
execution_mode: ExecutionMode | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
activities: ProtectionScopeActivities | str | int | Sequence[str] | None = None,
|
||||
locations: list[PolicyLocation | MutableMapping[str, Any]] | None = None,
|
||||
policy_actions: list[DlpActionInfo | MutableMapping[str, Any]] | None = None,
|
||||
execution_mode: ExecutionMode | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
# Extract aliased values from kwargs
|
||||
if "policyActions" in kwargs:
|
||||
policy_actions = kwargs["policyActions"]
|
||||
if "executionMode" in kwargs:
|
||||
execution_mode = kwargs["executionMode"]
|
||||
|
||||
# Deserialize activities flag
|
||||
if not isinstance(activities, ProtectionScopeActivities) and activities is not None:
|
||||
activities = deserialize_flag(activities, _PROTECTION_SCOPE_ACTIVITIES_MAP, ProtectionScopeActivities)
|
||||
|
||||
# Convert nested objects
|
||||
converted_locations: list[PolicyLocation] | None = None
|
||||
if locations is not None:
|
||||
converted_locations = cast(
|
||||
list[PolicyLocation],
|
||||
[loc if isinstance(loc, PolicyLocation) else PolicyLocation(**loc) for loc in locations],
|
||||
)
|
||||
|
||||
converted_policy_actions: list[DlpActionInfo] | None = None
|
||||
if policy_actions is not None:
|
||||
converted_policy_actions = cast(
|
||||
list[DlpActionInfo],
|
||||
[p if isinstance(p, DlpActionInfo) else DlpActionInfo(**p) for p in policy_actions],
|
||||
)
|
||||
|
||||
# Call parent without explicit params with aliases
|
||||
super().__init__(**kwargs)
|
||||
self.activities = activities # type: ignore[assignment]
|
||||
self.locations = converted_locations
|
||||
self.policy_actions = converted_policy_actions
|
||||
self.execution_mode = execution_mode
|
||||
|
||||
def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: # type: ignore[override]
|
||||
# Get base dict (activities will be missing because Flag isn't JSON-serializable)
|
||||
base = super().to_dict(exclude=exclude, exclude_none=exclude_none)
|
||||
|
||||
# Manually serialize activities flag if present and not excluded
|
||||
if self.activities is not None or not exclude_none:
|
||||
if self.activities is not None:
|
||||
base["activities"] = serialize_flag(self.activities, _PROTECTION_SCOPE_ACTIVITIES_SERIALIZE_ORDER)
|
||||
elif not exclude_none:
|
||||
base["activities"] = None
|
||||
|
||||
return base
|
||||
|
||||
|
||||
class ProtectionScopesResponse(_AliasSerializable):
|
||||
_ALIASES: ClassVar[dict[str, str]] = {"scope_identifier": "scopeIdentifier", "scopes": "value"}
|
||||
|
||||
scope_identifier: str | None
|
||||
scopes: list[PolicyScope] | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
scope_identifier: str | None = None,
|
||||
scopes: list[PolicyScope | MutableMapping[str, Any]] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
# Extract aliased values from kwargs before they're normalized by parent
|
||||
if "scopeIdentifier" in kwargs:
|
||||
scope_identifier = kwargs["scopeIdentifier"]
|
||||
if "value" in kwargs:
|
||||
scopes = kwargs["value"]
|
||||
|
||||
converted_scopes: list[PolicyScope] | None = None
|
||||
if scopes is not None:
|
||||
converted_scopes = cast(
|
||||
list[PolicyScope], [s if isinstance(s, PolicyScope) else PolicyScope(**s) for s in scopes]
|
||||
)
|
||||
|
||||
# Don't pass parameters that have aliases - let parent normalize them
|
||||
super().__init__(**kwargs)
|
||||
self.scope_identifier = scope_identifier
|
||||
self.scopes = converted_scopes
|
||||
|
||||
|
||||
class ContentActivitiesResponse(_AliasSerializable):
|
||||
DEFAULT_EXCLUDE: ClassVar[set[str]] = {"status_code"}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
status_code: int | None = None,
|
||||
error: ErrorDetails | MutableMapping[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
if isinstance(error, MutableMapping):
|
||||
error = ErrorDetails(**error)
|
||||
super().__init__(status_code=status_code, error=error, **kwargs)
|
||||
self.status_code = status_code
|
||||
self.error = error # type: ignore[assignment]
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AccessedResourceDetails",
|
||||
"Activity",
|
||||
"ActivityMetadata",
|
||||
"AiAgentInfo",
|
||||
"AiInteractionPlugin",
|
||||
"ContentActivitiesRequest",
|
||||
"ContentActivitiesResponse",
|
||||
"ContentBase",
|
||||
"ContentToProcess",
|
||||
"DeviceMetadata",
|
||||
"DlpAction",
|
||||
"DlpActionInfo",
|
||||
"ExecutionMode",
|
||||
"GraphDataTypeBase",
|
||||
"IntegratedAppMetadata",
|
||||
"OperatingSystemSpecifications",
|
||||
"PolicyLocation",
|
||||
"PolicyPivotProperty",
|
||||
"PolicyScope",
|
||||
"ProcessContentRequest",
|
||||
"ProcessContentResponse",
|
||||
"ProcessConversationMetadata",
|
||||
"ProcessingError",
|
||||
"ProtectedAppMetadata",
|
||||
"ProtectionScopeActivities",
|
||||
"ProtectionScopeState",
|
||||
"ProtectionScopesRequest",
|
||||
"ProtectionScopesResponse",
|
||||
"PurviewBinaryContent",
|
||||
"PurviewTextContent",
|
||||
"RestrictionAction",
|
||||
"deserialize_flag",
|
||||
"serialize_flag",
|
||||
"translate_activity",
|
||||
]
|
||||
@@ -0,0 +1,251 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from collections.abc import Iterable, MutableMapping
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import ChatMessage
|
||||
|
||||
from ._client import PurviewClient
|
||||
from ._models import (
|
||||
Activity,
|
||||
ActivityMetadata,
|
||||
ContentActivitiesRequest,
|
||||
ContentToProcess,
|
||||
DeviceMetadata,
|
||||
DlpAction,
|
||||
DlpActionInfo,
|
||||
IntegratedAppMetadata,
|
||||
OperatingSystemSpecifications,
|
||||
PolicyLocation,
|
||||
ProcessContentRequest,
|
||||
ProcessContentResponse,
|
||||
ProcessConversationMetadata,
|
||||
ProcessingError,
|
||||
ProtectedAppMetadata,
|
||||
ProtectionScopesRequest,
|
||||
ProtectionScopesResponse,
|
||||
PurviewTextContent,
|
||||
RestrictionAction,
|
||||
translate_activity,
|
||||
)
|
||||
from ._settings import PurviewSettings
|
||||
|
||||
|
||||
def _is_valid_guid(value: str | None) -> bool:
|
||||
"""Check if a string is a valid GUID/UUID format using uuid module."""
|
||||
if not value:
|
||||
return False
|
||||
try:
|
||||
uuid.UUID(value)
|
||||
return True
|
||||
except (ValueError, AttributeError):
|
||||
return False
|
||||
|
||||
|
||||
class ScopedContentProcessor:
|
||||
"""Combine protection scopes, process content, and content activities logic."""
|
||||
|
||||
def __init__(self, client: PurviewClient, settings: PurviewSettings):
|
||||
self._client = client
|
||||
self._settings = settings
|
||||
|
||||
async def process_messages(
|
||||
self, messages: Iterable[ChatMessage], activity: Activity, user_id: str | None = None
|
||||
) -> tuple[bool, str | None]:
|
||||
"""Process messages for policy evaluation.
|
||||
|
||||
Args:
|
||||
messages: The messages to process
|
||||
activity: The activity type (e.g., UPLOAD_TEXT)
|
||||
user_id: Optional user_id to use for all messages. If provided, this is the fallback.
|
||||
|
||||
Returns:
|
||||
A tuple of (should_block: bool, resolved_user_id: str | None).
|
||||
The resolved_user_id can be stored and passed back when processing the response
|
||||
to ensure the same user context is maintained throughout the request/response cycle.
|
||||
"""
|
||||
pc_requests, resolved_user_id = await self._map_messages(messages, activity, user_id)
|
||||
should_block = False
|
||||
for req in pc_requests:
|
||||
resp = await self._process_with_scopes(req)
|
||||
if resp.policy_actions:
|
||||
for act in resp.policy_actions:
|
||||
if act.action == DlpAction.BLOCK_ACCESS or act.restriction_action == RestrictionAction.BLOCK:
|
||||
should_block = True
|
||||
break
|
||||
if should_block:
|
||||
break
|
||||
return should_block, resolved_user_id
|
||||
|
||||
async def _map_messages(
|
||||
self, messages: Iterable[ChatMessage], activity: Activity, provided_user_id: str | None = None
|
||||
) -> tuple[list[ProcessContentRequest], str | None]:
|
||||
"""Map messages to ProcessContentRequests.
|
||||
|
||||
Args:
|
||||
messages: The messages to map
|
||||
activity: The activity type
|
||||
provided_user_id: Optional user_id to use. If provided, this is the fallback.
|
||||
|
||||
Returns:
|
||||
A tuple of (requests, resolved_user_id)
|
||||
"""
|
||||
results: list[ProcessContentRequest] = []
|
||||
token_info = None
|
||||
|
||||
if not (self._settings.tenant_id and self._settings.purview_app_location):
|
||||
token_info = await self._client.get_user_info_from_token(tenant_id=self._settings.tenant_id)
|
||||
|
||||
tenant_id = (token_info or {}).get("tenant_id") or self._settings.tenant_id
|
||||
if not tenant_id or not _is_valid_guid(tenant_id):
|
||||
raise ValueError("Tenant id required or must be inferable from credential")
|
||||
|
||||
resolved_user_id = (token_info or {}).get("user_id")
|
||||
resolved_author_name = None
|
||||
if not resolved_user_id:
|
||||
for m in messages:
|
||||
if m.additional_properties:
|
||||
potential_user_id = m.additional_properties.get("user_id")
|
||||
if _is_valid_guid(potential_user_id):
|
||||
resolved_user_id = potential_user_id
|
||||
break
|
||||
if m.author_name and _is_valid_guid(m.author_name) and not resolved_author_name:
|
||||
resolved_author_name = m.author_name
|
||||
|
||||
if not resolved_user_id and resolved_author_name:
|
||||
resolved_user_id = resolved_author_name
|
||||
|
||||
if not resolved_user_id:
|
||||
resolved_user_id = provided_user_id if provided_user_id and _is_valid_guid(provided_user_id) else None
|
||||
|
||||
# Return empty results if user_id is empty
|
||||
if not resolved_user_id or not _is_valid_guid(resolved_user_id):
|
||||
return results, None
|
||||
|
||||
for m in messages:
|
||||
message_id = m.message_id or str(uuid.uuid4())
|
||||
content = PurviewTextContent(data=m.text or "")
|
||||
meta = ProcessConversationMetadata(
|
||||
identifier=message_id,
|
||||
content=content,
|
||||
name=f"Agent Framework Message {message_id}",
|
||||
is_truncated=False,
|
||||
correlation_id=str(uuid.uuid4()),
|
||||
)
|
||||
activity_meta = ActivityMetadata(activity=activity)
|
||||
|
||||
if self._settings.purview_app_location:
|
||||
policy_location = PolicyLocation(
|
||||
data_type=self._settings.purview_app_location.get_policy_location()["@odata.type"],
|
||||
value=self._settings.purview_app_location.location_value,
|
||||
)
|
||||
elif token_info and token_info.get("client_id"):
|
||||
policy_location = PolicyLocation(
|
||||
data_type="microsoft.graph.policyLocationApplication",
|
||||
value=token_info["client_id"],
|
||||
)
|
||||
else:
|
||||
raise ValueError("App location not provided or inferable")
|
||||
|
||||
protected_app = ProtectedAppMetadata(
|
||||
name=self._settings.app_name,
|
||||
version="1.0",
|
||||
application_location=policy_location,
|
||||
)
|
||||
integrated_app = IntegratedAppMetadata(name=self._settings.app_name, version="1.0")
|
||||
device_meta = DeviceMetadata(
|
||||
operating_system_specifications=OperatingSystemSpecifications(
|
||||
operating_system_platform="Unknown", operating_system_version="Unknown"
|
||||
)
|
||||
)
|
||||
|
||||
ctp = ContentToProcess(
|
||||
content_entries=[meta],
|
||||
activity_metadata=activity_meta,
|
||||
device_metadata=device_meta,
|
||||
integrated_app_metadata=integrated_app,
|
||||
protected_app_metadata=protected_app,
|
||||
)
|
||||
req = ProcessContentRequest(
|
||||
content_to_process=ctp,
|
||||
user_id=resolved_user_id, # Use the resolved user_id for all messages
|
||||
tenant_id=tenant_id,
|
||||
correlation_id=meta.correlation_id,
|
||||
process_inline=True if self._settings.process_inline else None,
|
||||
)
|
||||
results.append(req)
|
||||
return results, resolved_user_id
|
||||
|
||||
async def _process_with_scopes(self, pc_request: ProcessContentRequest) -> ProcessContentResponse:
|
||||
app_location = pc_request.content_to_process.protected_app_metadata.application_location
|
||||
locations: list[PolicyLocation | MutableMapping[str, Any]] = [app_location] if app_location is not None else []
|
||||
|
||||
ps_req = ProtectionScopesRequest(
|
||||
user_id=pc_request.user_id,
|
||||
tenant_id=pc_request.tenant_id,
|
||||
activities=translate_activity(pc_request.content_to_process.activity_metadata.activity),
|
||||
locations=locations,
|
||||
device_metadata=pc_request.content_to_process.device_metadata,
|
||||
integrated_app_metadata=pc_request.content_to_process.integrated_app_metadata,
|
||||
correlation_id=pc_request.correlation_id,
|
||||
)
|
||||
ps_resp = await self._client.get_protection_scopes(ps_req)
|
||||
should_process, dlp_actions = self._check_applicable_scopes(pc_request, ps_resp)
|
||||
|
||||
if should_process:
|
||||
pc_resp = await self._client.process_content(pc_request)
|
||||
pc_resp.policy_actions = self._combine_policy_actions(pc_resp.policy_actions, dlp_actions)
|
||||
return pc_resp
|
||||
ca_req = ContentActivitiesRequest(
|
||||
user_id=pc_request.user_id,
|
||||
tenant_id=pc_request.tenant_id,
|
||||
content_to_process=pc_request.content_to_process,
|
||||
correlation_id=pc_request.correlation_id,
|
||||
)
|
||||
ca_resp = await self._client.send_content_activities(ca_req)
|
||||
if ca_resp.error:
|
||||
return ProcessContentResponse(processing_errors=[ProcessingError(message=str(ca_resp.error))])
|
||||
return ProcessContentResponse()
|
||||
|
||||
@staticmethod
|
||||
def _combine_policy_actions(
|
||||
existing: list[DlpActionInfo] | None, new_actions: list[DlpActionInfo]
|
||||
) -> list[DlpActionInfo]:
|
||||
by_key: dict[str, DlpActionInfo] = {}
|
||||
for a in existing or []:
|
||||
if a.action:
|
||||
by_key[a.action] = a
|
||||
for a in new_actions:
|
||||
if a.action:
|
||||
by_key[a.action] = a
|
||||
return list(by_key.values())
|
||||
|
||||
@staticmethod
|
||||
def _check_applicable_scopes(
|
||||
pc_request: ProcessContentRequest, ps_response: ProtectionScopesResponse
|
||||
) -> tuple[bool, list[DlpActionInfo]]:
|
||||
req_activity = translate_activity(pc_request.content_to_process.activity_metadata.activity)
|
||||
location = pc_request.content_to_process.protected_app_metadata.application_location
|
||||
should_process: bool = False
|
||||
dlp_actions: list[DlpActionInfo] = []
|
||||
for scope in ps_response.scopes or []:
|
||||
# Check if all activities in req_activity are present in scope.activities using bitwise flags.
|
||||
activity_match = bool(scope.activities and (scope.activities & req_activity) == req_activity)
|
||||
location_match = False
|
||||
if location is not None:
|
||||
for loc in scope.locations or []:
|
||||
if (
|
||||
loc.data_type
|
||||
and location.data_type
|
||||
and loc.data_type.lower().endswith(location.data_type.split(".")[-1].lower())
|
||||
and loc.value == location.value
|
||||
):
|
||||
location_match = True
|
||||
break
|
||||
if activity_match and location_match:
|
||||
should_process = True
|
||||
if scope.policy_actions:
|
||||
dlp_actions.extend(scope.policy_actions)
|
||||
return should_process, dlp_actions
|
||||
@@ -0,0 +1,71 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import Enum
|
||||
|
||||
from agent_framework._pydantic import AFBaseSettings
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic_settings import SettingsConfigDict
|
||||
|
||||
|
||||
class PurviewLocationType(str, Enum):
|
||||
"""The type of location for Purview policy evaluation."""
|
||||
|
||||
APPLICATION = "application"
|
||||
URI = "uri"
|
||||
DOMAIN = "domain"
|
||||
|
||||
|
||||
class PurviewAppLocation(BaseModel):
|
||||
"""Identifier representing the app's location for Purview policy evaluation."""
|
||||
|
||||
location_type: PurviewLocationType = Field(..., description="The location type.")
|
||||
location_value: str = Field(..., description="The location value.")
|
||||
|
||||
def get_policy_location(self) -> dict[str, str]:
|
||||
ns = "microsoft.graph"
|
||||
if self.location_type == PurviewLocationType.APPLICATION:
|
||||
dt = f"{ns}.policyLocationApplication"
|
||||
elif self.location_type == PurviewLocationType.URI:
|
||||
dt = f"{ns}.policyLocationUrl"
|
||||
elif self.location_type == PurviewLocationType.DOMAIN:
|
||||
dt = f"{ns}.policyLocationDomain"
|
||||
else: # pragma: no cover - defensive
|
||||
raise ValueError("Invalid Purview location type")
|
||||
return {"@odata.type": dt, "value": self.location_value}
|
||||
|
||||
|
||||
class PurviewSettings(AFBaseSettings):
|
||||
"""Settings for Purview integration mirroring .NET PurviewSettings.
|
||||
|
||||
Attributes:
|
||||
app_name: Public app name.
|
||||
tenant_id: Optional tenant id (guid) of the user making the request.
|
||||
purview_app_location: Optional app location for policy evaluation.
|
||||
graph_base_uri: Base URI for Microsoft Graph.
|
||||
blocked_prompt_message: Custom message to return when a prompt is blocked by policy.
|
||||
blocked_response_message: Custom message to return when a response is blocked by policy.
|
||||
"""
|
||||
|
||||
app_name: str = Field(...)
|
||||
tenant_id: str | None = Field(default=None)
|
||||
purview_app_location: PurviewAppLocation | None = Field(default=None)
|
||||
graph_base_uri: str = Field(default="https://graph.microsoft.com/v1.0/")
|
||||
process_inline: bool = Field(default=False, description="Process content inline if supported.")
|
||||
blocked_prompt_message: str = Field(
|
||||
default="Prompt blocked by policy",
|
||||
description="Message to return when a prompt is blocked by policy.",
|
||||
)
|
||||
blocked_response_message: str = Field(
|
||||
default="Response blocked by policy",
|
||||
description="Message to return when a response is blocked by policy.",
|
||||
)
|
||||
|
||||
model_config = SettingsConfigDict(populate_by_name=True, validate_assignment=True)
|
||||
|
||||
def get_scopes(self) -> list[str]:
|
||||
from urllib.parse import urlparse
|
||||
|
||||
host = urlparse(self.graph_base_uri).hostname or "graph.microsoft.com"
|
||||
return [f"https://{host}/.default"]
|
||||
Reference in New Issue
Block a user