agents/function_tools.py

"""Advisory function-call tools — PATH 1 of the agent harness.

This module holds the five advisory ``function_call`` actions (CVE lookup, risk
scoring, and drafting), exposed to the model as native Pydantic-AI tools
(``@agent.tool``). It is split out of :mod:`agents.soc_agent` so that file stays
focused on the loop, the brain, and the builders.

Each tool is a thin, typed wrapper around :func:`_run_action`, which fixes the
backend + action name and routes through the harness's shared governance helper
(:func:`agents.soc_agent.govern_action`) — the *same* checkpoint the structured-
draft loop uses, so an advisory tool call is governed exactly like any other
action. Typed tool parameters give the model a precise argument schema, which is
what makes native tool-calling reliable here.

Note the import direction: this module depends on :mod:`agents.soc_agent` for the
shared governance primitives, and ``soc_agent.build_agent`` imports
:func:`register_function_tools` lazily (inside the function) to register them. That
keeps the dependency one-way at import time, so there is no circular import.
"""

from __future__ import annotations

from typing import Any

from pydantic_ai import Agent, RunContext

from agents.soc_agent import GovernanceContext, _action_signature, govern_action
from control_plane.schema import Backend, ProposedAction


#: The advisory action names exposed as native tools. Kept as data so a test (and
#: a reader) can see the exact tool surface at a glance.
FUNCTION_CALL_TOOLS: tuple[str, ...] = (
    "cve_lookup",
    "risk_score",
    "draft_remediation",
    "draft_incident",
    "recommend_containment",
)


def _run_action(
    deps: GovernanceContext, action_name: str, arguments: dict[str, Any]
) -> dict[str, Any]:
    """Govern one advisory action and return the outcome to the model as a dict.

    The five tools below are thin typed fronts for this function. Each fixes
    ``backend = FUNCTION_CALL`` and its own ``action_name`` (the model can't forge
    them); the model supplies only the business arguments. Returns a small JSON
    object the model reads as feedback — never an exception.
    """
    arguments = {k: v for k, v in arguments.items() if v is not None}

    action = ProposedAction(
        incident_id=deps.incident_id,
        agent_id=deps.agent_id,
        # The model no longer chooses a tier; we record the operator ceiling the
        # gate evaluates against. The verdict stands on that ceiling alone.
        autonomy_tier=deps.operator_tier,
        backend=Backend.FUNCTION_CALL,
        action_name=action_name,
        arguments=arguments,
        reason=f"Agent requested {action_name} for incident {deps.incident_id}",
    )

    # Repeat nudge: a live model can loop on the same advisory call — re-running a
    # successful one, OR (just as wasteful) re-trying one that was already denied —
    # burning its step budget without making progress. If this exact action+args was
    # already attempted in this run, short-circuit with a firm "don't repeat" note
    # instead of governing a duplicate, steering the model to move on or conclude.
    signature = _action_signature(action)
    for prior in deps.turns:
        if _action_signature(prior.action) != signature:
            continue
        if prior.executed:
            # Already succeeded — hand back the same result, tell it not to repeat.
            return {
                "governance": "ALLOWED",
                "note": (
                    f"You already ran {action_name} with these arguments earlier in "
                    "this run; the result is unchanged. Do NOT call it again — use "
                    "the earlier result, take a different step, or give your FINAL "
                    "ANSWER."
                ),
                "result": prior.result,
            }
        # Already attempted and NOT allowed — repeating it will be denied again.
        return {
            "governance": prior.decision.value,
            "note": (
                f"You already tried {action_name} with these arguments and it was "
                f"not allowed ({prior.decision.value}). Do NOT repeat it — take a "
                "different, allowed step or give your FINAL ANSWER."
            ),
            "reason": prior.reason,
        }

    turn = govern_action(deps, action)
    if turn.executed:
        return {"governance": "ALLOWED", "result": turn.result}
    if turn.execution_error is not None:
        return {"governance": "ALLOWED", "executed": False, "error": turn.execution_error}
    return {"governance": turn.decision.value, "reason": turn.reason}


def register_function_tools(agent: Agent[GovernanceContext, Any]) -> None:
    """Register the five advisory `function_call` tools on *agent* (once, at build).

    Each tool is a thin, typed wrapper: the model sees a clear signature (e.g.
    ``cve_lookup(cve_id: str)``) and fills the arguments, while the body fixes the
    backend + action name and hands off to the shared governance helper. The typed
    parameters are deliberate — they give the model a precise argument schema,
    which is what makes native tool-calling reliable here.
    """

    @agent.tool
    def cve_lookup(ctx: RunContext[GovernanceContext], cve_id: str) -> dict[str, Any]:
        """Look up a CVE record (advisory, read-only)."""
        return _run_action(ctx.deps, "cve_lookup", {"cve_id": cve_id})

    @agent.tool
    def risk_score(
        ctx: RunContext[GovernanceContext], cve_id: str, asset_id: str | None = None
    ) -> dict[str, Any]:
        """Score a CVE's risk, optionally weighted by an asset's criticality (advisory)."""
        return _run_action(ctx.deps, "risk_score", {"cve_id": cve_id, "asset_id": asset_id})

    @agent.tool
    def draft_remediation(ctx: RunContext[GovernanceContext], cve_id: str) -> dict[str, Any]:
        """Draft remediation guidance for a CVE — text only, nothing is applied."""
        return _run_action(ctx.deps, "draft_remediation", {"cve_id": cve_id})

    @agent.tool
    def draft_incident(ctx: RunContext[GovernanceContext]) -> dict[str, Any]:
        """Draft an incident summary from this incident's alerts — text only."""
        return _run_action(ctx.deps, "draft_incident", {"incident_id": ctx.deps.incident_id})

    @agent.tool
    def recommend_containment(
        ctx: RunContext[GovernanceContext], asset_id: str, container_id: str | None = None
    ) -> dict[str, Any]:
        """Recommend containment steps for an asset — advisory only, nothing is done."""
        return _run_action(
            ctx.deps, "recommend_containment", {"asset_id": asset_id, "container_id": container_id}
        )