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."""
    try:
        result = _client.check(product, version)
        return {
            "outside_coverage": False,
            "risk_state": result.risk_state,
            "actively_exploited": result.actively_exploited,
            "patch_available": result.patch_available,
            "fixed_version": result.fixed_version,
            "supply_chain_compromised": (
                result.supply_chain.compromised
                if result.supply_chain is not None
                else False
            ),
        }
    except AttestdUnsupportedProductError:
        return {
            "outside_coverage": True,
            "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', or if supply_chain_compromised 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', or if "
        "supply_chain_compromised 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())