AutoGen Integration

pip install "autogen-agentchat==0.7.5" autogen-core autogen-ext[openai] attestd

import os

import attestd
from autogen_core.tools import FunctionTool
from attestd import AttestdUnsupportedProductError

# Instantiate once — captured in closure. Never instantiate inside the tool function.
_client = attestd.Client(api_key=os.environ["ATTESTD_API_KEY"])


def check_package_vulnerability(product: str, version: str) -> dict:
    """Check whether a software package version has known CVE vulnerabilities
    or supply chain compromise. Use before deploying or recommending any
    software dependency. outside_coverage=True means no data — treat as
    unknown risk, not safe. typosquat_detected=True means the name itself
    is suspect; use resembles to recover the intended package."""
    try:
        result = _client.check(product, version)
        sc = result.supply_chain
        return {
            "outside_coverage": False,
            "risk_state": result.risk_state,
            "remote_exploitable": result.remote_exploitable,
            "actively_exploited": result.actively_exploited,
            "patch_available": result.patch_available,
            "fixed_version": result.fixed_version,
            "supply_chain_compromised": sc.compromised if sc is not None else False,
            "supply_chain_malware_type": sc.malware_type if sc is not None else None,
            "supply_chain_still_live": (
                (sc.removed_at is None) if (sc is not None and sc.compromised) else None
            ),
        }
    except AttestdUnsupportedProductError as e:
        if e.typosquat and e.typosquat.detected:
            return {
                "outside_coverage": True,
                "typosquat_detected": True,
                "typosquat_kind": e.typosquat.kind,
                "resembles": e.typosquat.resembles,
                "likely_intended": list(e.typosquat.likely_intended),
                "message": (
                    f"'{product}' fails package name integrity "
                    f"(kind={e.typosquat.kind}). "
                    f"Likely intended: {e.typosquat.resembles}. "
                    f"Do not install this package."
                ),
            }
        return {
            "outside_coverage": True,
            "typosquat_detected": False,
            "risk_state": None,
            "message": f"No Attestd coverage for '{product}'. Treat as unknown risk.",
        }


attestd_tool = FunctionTool(
    check_package_vulnerability,
    description=(
        "Check CVE risk and supply chain integrity for a software dependency. "
        "outside_coverage=True means unknown risk, not safe."
    ),
)

import asyncio

from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.conditions import TextMentionTermination
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.ui import Console
from autogen_ext.models.openai import OpenAIChatCompletionClient

from attestd_tool import attestd_tool

model_client = OpenAIChatCompletionClient(model="gpt-4o-mini")

security_agent = AssistantAgent(
    name="security_agent",
    model_client=model_client,
    tools=[attestd_tool],
    system_message=(
        "You are a security-aware deployment assistant. "
        "Before approving any software dependency, call check_package_vulnerability. "
        "Block if risk_state is 'critical' or 'high', supply_chain_compromised is True, "
        "or typosquat_detected is True. "
        "If outside_coverage is True, state that explicitly. Do not treat it as safe. "
        "End with APPROVED or BLOCKED and your reasoning."
    ),
)

termination = TextMentionTermination("APPROVED", sources=["security_agent"]) | TextMentionTermination(
    "BLOCKED", sources=["security_agent"]
)
team = RoundRobinGroupChat([security_agent], termination_condition=termination)


async def main() -> None:
    await Console(
        team.run_stream(
            task="Is it safe to deploy with runc 1.0.0 and litellm 1.82.7?"
        )
    )


asyncio.run(main())

import asyncio

from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.conditions import TextMentionTermination
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.ui import Console
from autogen_ext.models.openai import OpenAIChatCompletionClient

from attestd_tool import attestd_tool

model_client = OpenAIChatCompletionClient(model="gpt-4o-mini")

security_gate = AssistantAgent(
    name="security_gate",
    model_client=model_client,
    tools=[attestd_tool],
    system_message=(
        "You are a security gate agent. Before any software dependency is "
        "deployed or recommended, check it with check_package_vulnerability. "
        "Block if risk_state is 'critical' or 'high', supply_chain_compromised is True, "
        "or typosquat_detected is True. "
        "If outside_coverage is True, state that explicitly. Do not treat it as safe. "
        "Respond with APPROVED or BLOCKED with your reasoning."
    ),
)

deployment_agent = AssistantAgent(
    name="deployment_agent",
    model_client=model_client,
    system_message=(
        "You handle deployment tasks. Wait for security_gate to approve "
        "all dependencies before proceeding."
    ),
)

termination = TextMentionTermination("APPROVED", sources=["security_gate"]) | TextMentionTermination(
    "BLOCKED", sources=["security_gate"]
)

team = RoundRobinGroupChat(
    [security_gate, deployment_agent],
    termination_condition=termination,
)


async def main() -> None:
    await Console(
        team.run_stream(
            task=(
                "Deploy a stack with runc 1.0.0, litellm 1.82.7, "
                "node-ipc 9.1.6, and nginx 1.27.4."
            )
        )
    )


asyncio.run(main())