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:
Rishabh Chawla
2025-10-16 14:46:04 -07:00
committed by GitHub
Unverified
parent 76ae0a62ac
commit 59da578902
24 changed files with 3707 additions and 5 deletions
@@ -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"]