Metadata-Version: 2.4
Name: maf-plainid
Version: 1.1.0
Summary: Microsoft Agent Framework integration for PlainID authorization
Author: PlainID
License-Expression: MIT
Classifier: Development Status :: 5 - Production/Stable
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Security
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: Typing :: Typed
Requires-Python: <3.14,>=3.10
Requires-Dist: agent-framework-core<2.0.0,>=1.6.0
Requires-Dist: core-plainid<3.0.0,>=2.0.0
Provides-Extra: all
Requires-Dist: core-plainid[all]<3.0.0,>=2.0.0; extra == 'all'
Requires-Dist: semantic-kernel<2.0.0,>=1.42.0; extra == 'all'
Provides-Extra: anonymization
Requires-Dist: core-plainid[anonymization]<3.0.0,>=2.0.0; extra == 'anonymization'
Provides-Extra: anonymization-ahds
Requires-Dist: core-plainid[anonymization-ahds]<3.0.0,>=2.0.0; extra == 'anonymization-ahds'
Provides-Extra: categorization-llm
Requires-Dist: core-plainid[categorization-llm]<3.0.0,>=2.0.0; extra == 'categorization-llm'
Provides-Extra: categorization-zeroshot
Requires-Dist: core-plainid[categorization-zeroshot]<3.0.0,>=2.0.0; extra == 'categorization-zeroshot'
Provides-Extra: retrieval
Requires-Dist: semantic-kernel<2.0.0,>=1.42.0; extra == 'retrieval'
Description-Content-Type: text/markdown

# maf-plainid

[PlainID](https://www.plainid.com/) authorization integration for the [Microsoft Agent Framework (MAF)](https://github.com/microsoft/agent-framework). Provides MAF middleware, context providers, and workflow executors for prompt categorization, text anonymization, and policy-based document retrieval across multiple vector stores.

This library depends on [core-plainid](https://pypi.org/project/core-plainid/) for the underlying authorization components (permissions provider, categorizer, anonymizer, PlainID clients, exceptions, etc.). Please refer to the **core-plainid README** for details on setting up those components.

All components are **async-only** and designed for use with MAF's async execution model.

## Installation

```bash
pip install maf-plainid
```

`core-plainid` and `agent-framework-core` are installed automatically as dependencies. Optional extras:

```bash
pip install maf-plainid[retrieval]                    # Semantic Kernel retrieval (MultiStoreRetriever, LambdaFilterProvider)
pip install maf-plainid[categorization-llm]           # LLM-based categorization via LiteLLM
pip install maf-plainid[categorization-zeroshot]      # Zero-shot classification via Hugging Face
pip install maf-plainid[anonymization]                # Presidio-based PII anonymization
pip install maf-plainid[anonymization-ahds]           # Anonymization + Azure Health De-identification
pip install maf-plainid[all]                          # Everything
```

> **Note:** The `retrieval` extra installs `semantic-kernel` which is required for `MultiStoreRetriever`, `LambdaFilterProvider`, and `RetrievalExecutor`. Categorization, anonymization, and SQL authorization work without it.

## Two Integration Styles

`maf-plainid` supports two integration approaches depending on your architecture:

| Style | Components | Best for |
|---|---|---|
| **Agent (middleware + context providers)** | `CategorizationMiddleware`, `AnonymizationMiddleware`, `RetrievalContextProvider` | Single-agent setups with `Agent` class |
| **Workflow (executors)** | `ContextSetupExecutor`, `CategorizationExecutor`, `RetrievalExecutor`, `AnonymizerExecutor`, `SQLAuthorizationExecutor` | Multi-step workflows with `WorkflowBuilder` |

Both styles read the `RequestContext` from MAF state and delegate to the same core-plainid components underneath.

## Passing Request Context

### Agent Style (Session State)

For middleware and context providers, store the `RequestContext` in the agent session state before calling `agent.run()`:

```python
from agent_framework import Agent, AgentSession
from core_plainid.models.context.request_context import AdditionalIdentity, RequestContext
from maf_plainid import PLAINID_REQUEST_CONTEXT_KEY

request_context = RequestContext(
    entity_id="your_entity_id",
    entity_type_id="your_entity_type",
    additional_identities=[
        AdditionalIdentity(
            entity_id="your_additional_entity_id",
            entity_type_id="your_additional_entity_type",
        ),
    ],
)

session = AgentSession()
session.state[PLAINID_REQUEST_CONTEXT_KEY] = request_context

response = await agent.run("What is the salary?", session=session)
```

### Workflow Style (ContextSetupExecutor)

For workflows, use `ContextSetupExecutor` as the start executor. Build a fresh workflow per request so each execution has isolated state:

```python
from core_plainid.models.context.request_context import RequestContext
from maf_plainid import ContextSetupExecutor

def create_workflow(request_context: RequestContext, ...):
    setup = ContextSetupExecutor(request_context=request_context)
    # ... add other executors and edges ...
    return WorkflowBuilder(start_executor=setup, ...).build()

workflow = create_workflow(request_context=ctx, ...)
result = await workflow.run("What is the salary?")
```

Multiple identities (e.g. a User and an AI Agent) are supported for agentic scenarios through the `additional_identities` field in `RequestContext`. Identity can also be resolved via HTTP headers — see the **Identity Context** section in the core-plainid README.

All core-plainid components support three authentication modes: client credentials, per-request JWT token, and automatic IDP token management via `IdpAuthProvider` — see the **Authentication** section in the core-plainid README.

## Categorization

### Middleware

The `CategorizationMiddleware` classifies the user's prompt against PlainID policies **before** the agent executes. If the categories are not allowed, a `PlainIDCategorizerException` is raised and the agent invocation is blocked.

For setting up the categorizer, classifier providers, and the PlainID `Prompt_Control` ruleset, see the **Category Filtering** section in the core-plainid README.

```python
from agent_framework import Agent
from core_plainid.categorization.categorizer import Categorizer
from core_plainid.utils.plainid_permissions_provider import PlainIDPermissionsProvider
from maf_plainid import CategorizationMiddleware

permissions_provider = PlainIDPermissionsProvider(
    base_url="https://platform-product.us1.plainid.io",
    client_id="your_client_id",
    client_secret="your_client_secret",
)

categorizer = Categorizer(
    classifier_provider=classifier,
    permissions_provider=permissions_provider,
    all_categories=["contract", "HR", "finance"],
)

agent = Agent(
    client=your_llm_client,
    name="my-agent",
    middleware=[CategorizationMiddleware(categorizer=categorizer)],
)

response = await agent.run("What is the weather today?", session=session)
```

### Executor

The `CategorizationExecutor` performs the same categorization check within a MAF workflow:

```python
from maf_plainid import CategorizationExecutor, ContextSetupExecutor

setup = ContextSetupExecutor(request_context=request_context)
categorizer_exec = CategorizationExecutor(
    categorizer=categorizer,
    next_executor_on_error="my_error_handler",  # optional
)

workflow = (
    WorkflowBuilder(start_executor=setup, output_from=[output])
    .add_edge(setup, categorizer_exec)
    .add_edge(categorizer_exec, output)
    .build()
)

result = await workflow.run("What is the weather today?")
```

## Anonymization

### Middleware

The `AnonymizationMiddleware` anonymizes the agent's **response** (post-execution). It runs after the agent produces a response and applies PII masking/encryption to the last assistant message.

For setting up the anonymizer, encryption key, AHDS, and the PlainID `Output_Control` ruleset, see the **Anonymization** section in the core-plainid README.

```python
from core_plainid.anonymization.presidio_anonymizer import PresidioAnonymizer
from core_plainid.utils.plainid_permissions_provider import PlainIDPermissionsProvider
from maf_plainid import AnonymizationMiddleware

permissions_provider = PlainIDPermissionsProvider(
    base_url="https://platform-product.us1.plainid.io",
    client_id="your_client_id",
    client_secret="your_client_secret",
)

anonymizer = PresidioAnonymizer(
    permissions_provider=permissions_provider,
    encrypt_key="your_16_char_key!",
)

agent = Agent(
    client=your_llm_client,
    name="my-agent",
    middleware=[AnonymizationMiddleware(anonymizer=anonymizer)],
)

response = await agent.run("Tell me about John Smith", session=session)
print(response.text)  # "*** is a senior engineer..."
```

### Executor

The `AnonymizerExecutor` anonymizes the input message within a workflow. By default it forwards the anonymized text as the outgoing message (`forward_result=True`) and stores it in workflow state:

```python
from maf_plainid import AnonymizerExecutor, ContextSetupExecutor

setup = ContextSetupExecutor(request_context=request_context)
anonymizer_exec = AnonymizerExecutor(
    anonymizer=anonymizer,
    next_executor_on_error="my_error_handler",  # optional
)

workflow = (
    WorkflowBuilder(start_executor=setup, output_from=[output])
    .add_edge(setup, anonymizer_exec)
    .add_edge(anonymizer_exec, output)
    .build()
)

result = await workflow.run("John Smith lives in New York")
# result.get_outputs() -> ["*** lives in ***"]
```

## SQL Authorization

The `SQLAuthorizationExecutor` authorizes SQL queries via the PlainID SQL Authorizer, applying row-level and column-level security. It takes the incoming message as the SQL query, authorizes it, stores the full response in workflow state, and by default forwards the modified SQL as the outgoing message.

```python
from core_plainid.clients.plainid_sql_authorizer_client import PlainIDSQLAuthorizerClient
from core_plainid.models.sql_authorization.sql_authorizer_options import (
    SQLAuthorizerFlagsOptions,
    SQLAuthorizerOptions,
)
from maf_plainid import ContextSetupExecutor, SQLAuthorizationExecutor

sql_client = PlainIDSQLAuthorizerClient(
    base_url="https://platform-product.us1.plainid.io",
    client_id="your_client_id",
    client_secret="your_client_secret",
)

options = SQLAuthorizerOptions(
    flags=SQLAuthorizerFlagsOptions(expand_star_column=True),
)

setup = ContextSetupExecutor(request_context=request_context)
sql_exec = SQLAuthorizationExecutor(
    sql_authorizer_client=sql_client,
    options=options,
    forward_result=True,  # forward modified SQL (default)
    next_executor_on_error="my_error_handler",  # optional
)

workflow = (
    WorkflowBuilder(start_executor=setup, output_from=[output])
    .add_edge(setup, sql_exec)
    .add_edge(sql_exec, output)
    .build()
)

result = await workflow.run("SELECT * FROM employees")
# result.get_outputs() -> ["SELECT name FROM employees WHERE dept = 'eng'"]
```

The full `SQLAuthorizerResponse` is stored in workflow state at `PLAINID_SQL_AUTHORIZATION_RESULT_KEY`, which includes `sql`, `was_modified`, and optional `error`.

### Limitation

The SQL Authorizer does not currently support the runtime interface for primary and secondary identities (User + Agent). Only a single identity context is supported.

## Retrieval

The retrieval system enforces PlainID authorization policies on document retrieval from vector stores. It supports **multiple vector stores** simultaneously, where each PlainID resource type maps to a single vector store collection (e.g. a Chroma collection, a FAISS index, or an Azure AI Search index).

### PlainID Setup

Configure rulesets in PlainID using a custom template name (one resource type per vector store collection). For example, if you have a `customer` collection with `country` and `age` metadata:

```
# METADATA
# custom:
#   plainid:
#     kind: Ruleset
#     name: rs1
ruleset(asset, identity, requestParams, action) if {
    asset.template == "customer"
    asset["country"] == "Sweden"
    asset["country"] != "Russia"
    asset["age"] >= 5
}
```

Note that you need to add `country` and `age` parameters to your Semantic Kernel data model as fields marked with `is_filterable=True`. PlainID uses these metadata fields to build the lambda filters applied during retrieval.

### LambdaFilterProvider

The `LambdaFilterProvider` connects to PlainID and converts resolution data into string lambda expressions compatible with Semantic Kernel's `VectorSearch.search(filter=...)` parameter. SK parses these lambdas as AST and translates them into the native filter format for each connector (Chroma where-dicts, FAISS NumPy filters, Azure AI Search OData, etc.).

```python
from maf_plainid.retrieval.lambda_filter_provider import LambdaFilterProvider

filter_provider = LambdaFilterProvider(
    base_url="https://platform-product.us1.plainid.io",
    client_id="your_client_id",
    client_secret="your_client_secret",
)
```

Generated filters look like:

```python
"lambda record: record.title == 'The Famous Cat'"
"lambda record: (record.country == 'Sweden') and (record.age >= 5)"
"lambda record: (record.title == 'Book A') or (record.title == 'Book B')"
```

### MultiStoreRetriever

The `MultiStoreRetriever` queries multiple Semantic Kernel `VectorSearch` instances concurrently, applying PlainID authorization filters to each store.

```python
from semantic_kernel.connectors.memory.chroma import ChromaCollection
from maf_plainid.retrieval.lambda_filter_provider import LambdaFilterProvider
from maf_plainid.retrieval.multi_store_retriever import MultiStoreRetriever

filter_provider = LambdaFilterProvider(
    base_url="https://platform-product.us1.plainid.io",
    client_id="your_client_id",
    client_secret="your_client_secret",
)

customer_store = ChromaCollection(
    collection_name="customers",
    record_type=CustomerRecord,
    embedding_generator=embeddings,
)
product_store = ChromaCollection(
    collection_name="products",
    record_type=ProductRecord,
    embedding_generator=embeddings,
)

retriever = MultiStoreRetriever(
    filter_provider=filter_provider,
    resource_types=["customer", "product"],
    vector_stores=[customer_store, product_store],
    k=4,
)

docs = await retriever.aretrieve(
    "What is the capital of Sweden?",
    request_context=request_context,
)
```

#### Parameters

| Parameter | Type | Required | Description |
|---|---|---|---|
| `filter_provider` | `LambdaFilterProvider` | Yes | Connects to PlainID and converts policies into SK lambda filters |
| `resource_types` | `list[str]` | Yes | PlainID resource types — one per vector store |
| `vector_stores` | `list[VectorSearch]` | Yes | SK VectorSearch instances, aligned by index with `resource_types` |
| `k` | `int` | No | Global max documents per store (default `4`) |
| `k_values` | `list[int]` | No | Per-store document limits, overrides `k` when provided |

The `resource_types` and `vector_stores` lists must be the same length — each resource type at index `i` maps to the vector store at the same index.

#### Authorization Behavior

- If a resource type is **present** in the PlainID resolution, documents from that store are retrieved with the corresponding filter applied.
- If a resource type is **absent** from the resolution, the entire store is skipped (no documents returned from it).
- If a resource type has an **unrestricted** policy (all access), documents are retrieved without filtering.

This means PlainID controls *which* stores are queried and *what* documents are visible, while the vector store's similarity search handles relevance ranking.

### Context Provider

The `RetrievalContextProvider` integrates retrieval into the agent-style flow as a MAF `ContextProvider`. It runs before the agent invocation, retrieves documents based on the latest user message, and stores them in session state:

```python
from maf_plainid import RetrievalContextProvider

provider = RetrievalContextProvider(retriever=retriever)

agent = Agent(
    client=your_llm_client,
    name="my-agent",
    context_providers=[provider],
)

response = await agent.run("What is the capital of Sweden?", session=session)
# Retrieved docs available at: session.state[PLAINID_RETRIEVED_DOCUMENTS_KEY]
```

### Executor

The `RetrievalExecutor` performs retrieval within a workflow:

```python
from maf_plainid import ContextSetupExecutor
from maf_plainid.executors.retrieval_executor import RetrievalExecutor

setup = ContextSetupExecutor(request_context=request_context)
retrieval_exec = RetrievalExecutor(
    retriever=retriever,
    forward_result=True,  # send docs as outgoing message (default: False)
    next_executor_on_error="my_error_handler",  # optional
)

workflow = (
    WorkflowBuilder(start_executor=setup, output_from=[output])
    .add_edge(setup, retrieval_exec)
    .add_edge(retrieval_exec, output)
    .build()
)

result = await workflow.run("What is the capital of Sweden?")
# When forward_result=False: result.get_outputs() -> ["What is the capital of Sweden?"]
# When forward_result=True: result.get_outputs() -> [[VectorSearchResult, ...]]
```

## Combining Components

### Agent Style (Middleware + Context Provider)

Middleware and context providers compose naturally on a single `Agent`:

```python
from maf_plainid import (
    AnonymizationMiddleware,
    CategorizationMiddleware,
    RetrievalContextProvider,
)

agent = Agent(
    client=your_llm_client,
    name="secure-agent",
    middleware=[
        CategorizationMiddleware(categorizer=categorizer),
        AnonymizationMiddleware(anonymizer=anonymizer),
    ],
    context_providers=[
        RetrievalContextProvider(retriever=retriever),
    ],
)

session = AgentSession()
session.state[PLAINID_REQUEST_CONTEXT_KEY] = request_context

response = await agent.run("What is John Smith's contract status?", session=session)
```

This agent will:

1. **Categorize** the prompt — verify it matches allowed categories (pre-execution middleware)
2. **Retrieve** documents — query vector stores with PlainID filters (context provider)
3. **Execute** the LLM call with retrieved context
4. **Anonymize** the response — mask PII in the assistant's reply (post-execution middleware)

### Workflow Style (Executors)

Executors are composed into a MAF workflow using `WorkflowBuilder`:

```python
from agent_framework import WorkflowBuilder
from maf_plainid import (
    AnonymizerExecutor,
    CategorizationExecutor,
    ContextSetupExecutor,
    RetrievalExecutor,
)

def create_workflow(request_context: RequestContext):
    setup = ContextSetupExecutor(request_context=request_context)
    categorizer_exec = CategorizationExecutor(
        categorizer=categorizer,
        next_executor_on_error="error_handler",
    )
    retrieval_exec = RetrievalExecutor(retriever=retriever)
    anonymizer_exec = AnonymizerExecutor(anonymizer=anonymizer)
    output = YourOutputExecutor(id="output")
    error_handler = YourErrorHandler(id="error_handler")

    return (
        WorkflowBuilder(start_executor=setup, output_from=[output])
        .add_edge(setup, categorizer_exec)
        .add_edge(categorizer_exec, retrieval_exec)
        .add_edge(categorizer_exec, error_handler)
        .add_edge(retrieval_exec, anonymizer_exec)
        .add_edge(retrieval_exec, error_handler)
        .add_edge(anonymizer_exec, output)
        .add_edge(anonymizer_exec, error_handler)
        .build()
    )

workflow = create_workflow(request_context=request_context)
result = await workflow.run("What is the salary?")
```

This workflow will:

1. **Setup** — inject `RequestContext` into workflow state
2. **Categorize** — verify the prompt against PlainID policies
3. **Retrieve** — query authorized vector stores
4. **Anonymize** — mask PII in the query text
5. **Output** — emit the final result

## Error Handling

### Middleware

When middleware throws (e.g. `PlainIDCategorizerException`), the exception propagates directly to the caller of `agent.run()`:

```python
from core_plainid.exceptions import PlainIDCategorizerException

try:
    response = await agent.run("blocked prompt", session=session)
except PlainIDCategorizerException as e:
    print(f"Prompt blocked: {e}")
```

### Workflow Executors

All PlainID executors extend `BaseExecutor`, which supports an optional `next_executor_on_error` parameter:

- **Without error handler**: exceptions propagate and the workflow fails.
- **With error handler**: exceptions are caught, wrapped in an `ErrorResult`, and sent to your specified error executor via `send_message(error, target_id=...)`.

Example error executor

```python
from agent_framework import Executor, WorkflowContext, handler
from maf_plainid import ErrorResult, PLAINID_ERROR_DETAILS_KEY

class MyErrorHandler(Executor):
    @handler
    async def handle_error(self, error: ErrorResult, ctx: WorkflowContext) -> None:
        # error.executor_id — which executor failed
        # error.error_message — the error message string
        # error.original_error — the original exception
        ctx.set_state(PLAINID_ERROR_DETAILS_KEY, error)
        # ... log, notify, return fallback response, etc.
```

**Important**: The error handler must have a direct edge from each executor that might route errors to it. MAF only delivers messages along declared edges:

```python
.add_edge(categorizer_exec, error_handler)   # required
.add_edge(retrieval_exec, error_handler)     # required
.add_edge(anonymizer_exec, error_handler)    # required
```

## Workflow State Keys

All executors read/write from workflow state using well-known keys:

| Key | Constant | Written by | Description |
|---|---|---|---|
| `plainid_request_context` | `PLAINID_REQUEST_CONTEXT_KEY` | `ContextSetupExecutor` | The `RequestContext` for the current request |
| `plainid_retrieved_documents` | `PLAINID_RETRIEVED_DOCUMENTS_KEY` | `RetrievalExecutor` | List of `VectorSearchResult` from retrieval |
| `plainid_anonymized_text` | `PLAINID_ANONYMIZED_TEXT_KEY` | `AnonymizerExecutor` | The anonymized text |
| `plainid_sql_authorization_result` | `PLAINID_SQL_AUTHORIZATION_RESULT_KEY` | `SQLAuthorizationExecutor` | `SQLAuthorizerResponse` with modified SQL |
| `plainid_error_details` | `PLAINID_ERROR_DETAILS_KEY` | `BaseExecutor` (on error) | `ErrorResult` with error details |

Within executors, state is accessed via the `WorkflowContext`:

```python
from maf_plainid import PLAINID_REQUEST_CONTEXT_KEY, PLAINID_SQL_AUTHORIZATION_RESULT_KEY

request_context = ctx.get_state(PLAINID_REQUEST_CONTEXT_KEY)
sql_result = ctx.get_state(PLAINID_SQL_AUTHORIZATION_RESULT_KEY)
```

## Supported Vector Stores

`maf-plainid` uses Semantic Kernel's `VectorSearch` abstraction. Any SK vector store connector that supports `search(filter=...)` with string lambda filters is compatible. Tested vector stores:

### Chroma

Supports: `EQUALS`, `NOTEQUALS`, `GREATE`, `LESS`, `GREAT_EQUALS`, `LESS_EQUALS`.

Does not support: `IN`, `NOT_IN`, `STARTSWITH`, `ENDSWITH`, `CONTAINS`.

```python
from semantic_kernel.connectors.memory.chroma import ChromaCollection

store = ChromaCollection(
    collection_name="my_collection",
    record_type=MyRecord,
    embedding_generator=embeddings,
)
```

### FAISS

Supports: `EQUALS`, `NOTEQUALS`, `GREATE`, `LESS`, `GREAT_EQUALS`, `LESS_EQUALS`, `IN`, `NOT_IN`.

Does not support: `STARTSWITH`, `ENDSWITH`, `CONTAINS`.

```python
from semantic_kernel.connectors.memory.faiss import FaissCollection

store = FaissCollection(
    collection_name="my_collection",
    record_type=MyRecord,
    embedding_generator=embeddings,
)
```

### Azure AI Search

Supports all operators. Requires `azure-search-documents` package.

```python
from semantic_kernel.connectors.memory.azure_ai_search import AzureAISearchCollection

store = AzureAISearchCollection(
    collection_name="my-index",
    record_type=MyRecord,
    search_endpoint="https://your-service.search.windows.net",
    api_key="your-api-key",
)
```

### Azure Cosmos DB

Supports all operators. Requires `azure-cosmos` package.

```python
from semantic_kernel.connectors.memory.azure_cosmos_db import AzureCosmosDBNoSQLCollection

store = AzureCosmosDBNoSQLCollection(
    collection_name="my_container",
    record_type=MyRecord,
    database_name="my_database",
    url="https://your-account.documents.azure.com:443/",
    key="your-key",
)
```

### Defining a Data Model

Semantic Kernel requires a data model class decorated with `@vectorstoremodel`:

```python
from dataclasses import dataclass
from semantic_kernel.data import vectorstoremodel, VectorStoreField

@vectorstoremodel
@dataclass
class CustomerRecord:
    id: str = VectorStoreField(is_key=True)
    content: str = VectorStoreField()
    embedding: list[float] = VectorStoreField(dimensions=1536, embedding_property_name="content")
    country: str = VectorStoreField(is_filterable=True)
    age: int = VectorStoreField(is_filterable=True)
```

Fields used in PlainID rulesets must have `is_filterable=True`.

## Exceptions

All exceptions are defined in the `core-plainid` library. See the **Exceptions** section in the core-plainid README for the full list.
