Feat/asqav audit plugin

.github/workflows/pythonbuild.yml

@@ -291,6 +291,7 @@ jobs:
         plugin-names:
           # Please maintain an alphabetical order in the following list
           - flytekit-airflow
+          - flytekit-asqav
           - flytekit-async-fsspec
           - flytekit-aws-athena
           - flytekit-aws-batch

plugins/flytekit-asqav/README.md

@@ -0,0 +1,66 @@
+# Flytekit Asqav Plugin
+
+Cryptographic audit trails for Flyte tasks using [Asqav](https://asqav.com).
+
+Each task execution is wrapped with cryptographically signed receipts at lifecycle
+points (started, finished, failed), providing provable, tamper-proof records for
+AI workflows in regulated environments (finance, healthcare, etc.).
+
+## Installation
+
+```bash
+pip install flytekitplugins-asqav
+```
+
+Requires `flytekit>=1.12.0` and `asqav>=0.1.0`.
+
+## Quick Start
+
+```python
+from flytekit import task, Secret, workflow
+from flytekitplugins.asqav import asqav_audit
+
+asqav_secret = Secret(key="asqav-api-key", group="asqav")
+
+@task(secret_requests=[asqav_secret])
+@asqav_audit(agent_name="model-trainer", secret=asqav_secret)
+def train_model() -> float:
+    # ... training code ...
+    return 0.95
+
+@workflow
+def main() -> float:
+    return train_model()
+```
+
+## Configuration
+
+### `@asqav_audit(...)`
+
+| Parameter | Type | Required | Default | Description |
+|---|---|---|---|---|
+| `agent_name` | `str` | ✅ | — | Asqav agent name for this task |
+| `secret` | `Secret` or `Callable` | ❌ | `ASQAV_API_KEY` env var | How to resolve the API key |
+| `receipt_type` | `str` | ❌ | `protectmcp:lifecycle` | Asqav receipt type |
+| `risk_class` | `str` | ❌ | `None` | Risk classification (`low`, `medium`, `high`, `unknown`) |
+| `compliance_mode` | `bool` | ❌ | `True` | Generate IETF-compliant compliance receipts |
+| `fail_closed` | `bool` | ❌ | `False` | If `True`, Asqav errors propagate (fail-closed); if `False`, errors are logged and the task continues (fail-open) |
+
+### Secret Resolution (priority order)
+
+1. **`Secret` object** — resolved via Flyte's `SecretsManager` (for remote deployment)
+2. **Callable** — zero-argument function returning the API key string
+3. **`ASQAV_API_KEY`** environment variable — fallback
+
+## Audit Trail Output
+
+Each receipt is logged and rendered as a Flyte Deck card with:
+- Receipt type (started / finished / failed)
+- Execution ID and timestamp
+- Duration (for finished/failed receipts)
+- Error info (for failed receipts)
+- Clickable verification link to Asqav
+
+## License
+
+Apache 2.0

plugins/flytekit-asqav/flytekitplugins/asqav/__init__.py

@@ -0,0 +1,19 @@
+"""
+.. currentmodule:: flytekitplugins.asqav
+
+This plugin provides cryptographic audit trails for Flyte tasks using
+`Asqav <https://asqav.com>`_. Each task execution is wrapped with
+cryptographically signed receipts at lifecycle points (started, finished,
+failed), enabling provable, tamper-proof audit records for AI workflows in
+regulated environments.
+
+.. autosummary::
+   :template: custom.rst
+   :toctree: generated/
+
+   asqav_audit
+"""
+
+from .tracking import asqav_audit
+
+__all__ = ["asqav_audit"]

plugins/flytekit-asqav/flytekitplugins/asqav/tracking.py

@@ -0,0 +1,257 @@
+import datetime
+import os
+from typing import Any, Callable, Dict, Optional, Union
+
+import asqav
+from flytekit import Deck, Secret
+from flytekit.core.context_manager import FlyteContextManager
+from flytekit.core.utils import ClassDecorator
+from flytekit.loggers import logger
+
+ASQAV_LINK_TYPE_KEY = "link_type"
+ASQAV_DEFAULT_RECEIPT_TYPE = "protectmcp:lifecycle"
+ASQAV_ENV_API_KEY = "ASQAV_API_KEY"
+
+ASQAV_SCOPE_KEY = "scope"
+ASQAV_AGENT_NAME_KEY = "agent_name"
+ASQAV_RECEIPT_TYPE_KEY = "receipt_type"
+ASQAV_RISK_CLASS_KEY = "risk_class"
+ASQAV_COMPLIANCE_MODE_KEY = "compliance_mode"
+
+
+def _receipt_html(receipt_type: str, sig_id: str, verification_url: str, metadata: Dict[str, Any]) -> str:
+    """Render a single audit receipt as an HTML row."""
+    status_icon = "✅" if receipt_type == "finished" else "❌" if receipt_type == "failed" else "⏳"
+    rows = "".join(
+        f"<tr><td style='padding: 4px 8px; font-weight: 600;'>{k}</td>"
+        f"<td style='padding: 4px 8px;'>{v}</td></tr>"
+        for k, v in metadata.items()
+    )
+    return f"""
+    <div style="border: 1px solid #e0e0e0; border-radius: 8px; margin-bottom: 16px; overflow: hidden;">
+        <div style="background: #f8f9fa; padding: 12px 16px; border-bottom: 1px solid #e0e0e0;
+                    font-size: 16px; font-weight: 600;">
+            {status_icon} Receipt: {receipt_type}
+            <span style="float: right; font-size: 12px; color: #666;">ID: {sig_id}</span>
+        </div>
+        <div style="padding: 12px 16px;">
+            <table style="width: 100%; border-collapse: collapse;">
+                {rows}
+            </table>
+            <p style="margin-top: 12px;">
+                <a href="{verification_url}" target="_blank"
+                   style="background: #4f46e5; color: white; padding: 8px 16px; border-radius: 4px;
+                          text-decoration: none; font-size: 13px;">
+                    🔗 Verify on Asqav
+                </a>
+            </p>
+        </div>
+    </div>"""
+
+
+class asqav_audit(ClassDecorator):
+    """Cryptographic audit decorator for Flyte tasks using Asqav.
+
+    Wraps a ``@task``-decorated function to sign cryptographically verifiable
+    receipts at each lifecycle point (started, finished, failed). Receipts are
+    rendered as a Flyte Deck card with verification links.
+
+    Usage::
+
+        from flytekit import task, Secret, workflow
+        from flytekitplugins.asqav import asqav_audit
+
+        asqav_secret = Secret(key="asqav-api-key", group="asqav")
+
+        @task(secret_requests=[asqav_secret])
+        @asqav_audit(agent_name="model-trainer", secret=asqav_secret)
+        def train_model() -> float:
+            ...
+    """
+
+    def __init__(
+        self,
+        task_function=None,
+        *,
+        agent_name: str,
+        secret: Optional[Union[Secret, Callable]] = None,
+        receipt_type: str = ASQAV_DEFAULT_RECEIPT_TYPE,
+        risk_class: Optional[str] = None,
+        compliance_mode: bool = True,
+        fail_closed: bool = False,
+        **init_kwargs,
+    ):
+        if not agent_name:
+            raise ValueError("agent_name must be set")
+        super().__init__(
+            task_function,
+            agent_name=agent_name,
+            secret=secret,
+            receipt_type=receipt_type,
+            risk_class=risk_class,
+            compliance_mode=compliance_mode,
+            fail_closed=fail_closed,
+            **init_kwargs,
+        )
+        self.agent_name = agent_name
+        self.secret = secret
+        self.receipt_type = receipt_type
+        self.risk_class = risk_class
+        self.compliance_mode = compliance_mode
+        self.fail_closed = fail_closed
+        self._agent = None
+
+    def _resolve_api_key(self) -> str:
+        """Resolve the Asqav API key from the most controlled source available."""
+        # 1. Flyte Secret — highest priority for remote deployment
+        if isinstance(self.secret, Secret):
+            ctx = FlyteContextManager.current_context()
+            try:
+                return ctx.user_space_params.secrets.get(key=self.secret.key, group=self.secret.group)
+            except ValueError:
+                pass  # Secret not available in this context (e.g. local execution), fall through
+        # 2. Callable — for programmatic injection
+        if callable(self.secret):
+            return self.secret()
+        # 3. Environment variable — fallback for local development
+        api_key = os.environ.get(ASQAV_ENV_API_KEY)
+        if api_key:
+            return api_key
+        raise ValueError(
+            f"No Asqav API key found. Provide a `Secret` or callable to "
+            f"`asqav_audit(secret=...)`, or set the {ASQAV_ENV_API_KEY} env var."
+        )
+
+    def get_extra_config(self):
+        extra_config = {
+            ASQAV_LINK_TYPE_KEY: "asqav-audit",
+            ASQAV_AGENT_NAME_KEY: self.agent_name,
+            ASQAV_SCOPE_KEY: "flyte-task-lifecycle",
+            ASQAV_RECEIPT_TYPE_KEY: self.receipt_type,
+            ASQAV_COMPLIANCE_MODE_KEY: str(self.compliance_mode),
+        }
+        if self.risk_class:
+            extra_config[ASQAV_RISK_CLASS_KEY] = self.risk_class
+        return extra_config
+
+    def execute(self, *args, **kwargs):
+        ctx = FlyteContextManager.current_context()
+        is_local = ctx.execution_state.is_local_execution()
+        execution_id = str(ctx.user_space_params.execution_id) if not is_local else "local"
+
+        api_key = self._resolve_api_key()
+
+        receipts = []
+        start_time = datetime.datetime.utcnow()
+
+        # --- init Asqav agent (cached per process) and sign started receipt ---
+        agent = None
+        try:
+            if self._agent is None:
+                asqav.init(api_key=api_key)
+                self._agent = asqav.Agent.create(name=self.agent_name)
+            agent = self._agent
+
+            started_context = {
+                "event_type": "flyte.task.started",
+                "execution_id": execution_id,
+                "agent_name": self.agent_name,
+                "timestamp": datetime.datetime.utcnow().isoformat() + "Z",
+            }
+            started_resp = agent.sign(
+                action_type="flyte.task.started",
+                context=started_context,
+                compliance_mode=self.compliance_mode,
+                receipt_type=self.receipt_type,
+                risk_class=self.risk_class,
+            )
+            receipts.append(("started", started_resp))
+            logger.info(f"Signed asqav started receipt: {started_resp.signature_id}")
+        except Exception as e:
+            if self.fail_closed:
+                raise
+            logger.error(f"Asqav audit trail init/started receipt failed: {e}")
+
+        # --- task execution ---
+        output = None
+        exception: Optional[Exception] = None
+        try:
+            output = self.task_function(*args, **kwargs)
+        except Exception as e:
+            exception = e
+            raise
+        finally:
+            if agent is not None:
+                elapsed = (datetime.datetime.utcnow() - start_time).total_seconds()
+
+                if exception is not None:
+                    error_info = f"{type(exception).__name__}: {exception}"
+                    failed_context = {
+                        "event_type": "flyte.task.failed",
+                        "execution_id": execution_id,
+                        "agent_name": self.agent_name,
+                        "duration_seconds": round(elapsed, 3),
+                        "error": error_info,
+                        "timestamp": datetime.datetime.utcnow().isoformat() + "Z",
+                    }
+                    try:
+                        failed_resp = agent.sign(
+                            action_type="flyte.task.failed",
+                            context=failed_context,
+                            compliance_mode=self.compliance_mode,
+                            receipt_type=self.receipt_type,
+                            risk_class=self.risk_class,
+                        )
+                        receipts.append(("failed", failed_resp))
+                        logger.info(f"Signed asqav failed receipt: {failed_resp.signature_id}")
+                    except Exception as sign_err:
+                        if self.fail_closed:
+                            raise
+                        logger.error(f"Failed to sign asqav receipt for error: {sign_err}")
+                else:
+                    finished_context = {
+                        "event_type": "flyte.task.finished",
+                        "execution_id": execution_id,
+                        "agent_name": self.agent_name,
+                        "duration_seconds": round(elapsed, 3),
+                        "output_type": type(output).__name__ if output is not None else "None",
+                        "timestamp": datetime.datetime.utcnow().isoformat() + "Z",
+                    }
+                    try:
+                        finished_resp = agent.sign(
+                            action_type="flyte.task.finished",
+                            context=finished_context,
+                            compliance_mode=self.compliance_mode,
+                            receipt_type=self.receipt_type,
+                            risk_class=self.risk_class,
+                        )
+                        receipts.append(("finished", finished_resp))
+                        logger.info(f"Signed asqav finished receipt: {finished_resp.signature_id}")
+                    except Exception as sign_err:
+                        if self.fail_closed:
+                            raise
+                        logger.error(f"Failed to sign asqav receipt for success: {sign_err}")
+
+            # Render Flyte Deck with all receipts
+            if receipts:
+                deck_html_parts = [
+                    "<h3>Asqav Cryptographic Audit Trail</h3>",
+                    f"<p><strong>Agent:</strong> {self.agent_name} | "
+                    f"<strong>Execution ID:</strong> {execution_id}</p>",
+                ]
+                for rtype, resp in receipts:
+                    metadata = {
+                        "Action Type": resp.signature_id,
+                        "Algorithm": resp.algorithm or "unknown",
+                        "Compliance Mode": str(resp.compliance_mode),
+                    }
+                    if resp.receipt_type:
+                        metadata["Receipt Type"] = resp.receipt_type
+                    if resp.risk_class:
+                        metadata["Risk Class"] = resp.risk_class
+                    deck_html_parts.append(
+                        _receipt_html(rtype, resp.signature_id, resp.verification_url, metadata)
+                    )
+                deck = Deck("Asqav Audit Trail", "\n".join(deck_html_parts))
+
+        return output

plugins/flytekit-asqav/setup.py

@@ -0,0 +1,37 @@
+from setuptools import setup
+
+PLUGIN_NAME = "asqav"
+
+microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
+
+plugin_requires = ["flytekit>=1.12.0", "asqav>=0.1.0"]
+
+__version__ = "0.0.0+develop"
+
+setup(
+    title="Asqav",
+    title_expanded="Flytekit Asqav Audit Trail Plugin",
+    name=microlib_name,
+    version=__version__,
+    author="flyteorg",
+    author_email="admin@flyte.org",
+    description="Cryptographic audit trails for Flyte tasks using Asqav",
+    namespace_packages=["flytekitplugins"],
+    packages=[f"flytekitplugins.{PLUGIN_NAME}"],
+    install_requires=plugin_requires,
+    license="apache2",
+    python_requires=">=3.10",
+    classifiers=[
+        "Intended Audience :: Science/Research",
+        "Intended Audience :: Developers",
+        "License :: OSI Approved :: Apache Software License",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+        "Topic :: Scientific/Engineering",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        "Topic :: Software Development",
+        "Topic :: Software Development :: Libraries",
+        "Topic :: Software Development :: Libraries :: Python Modules",
+    ],
+)