#!/usr/bin/env python3
"""
Vendor Risk Assessment Agent — runnable demo (Agent Blueprint from AgentKits)

This runs a REAL agent loop: it sends the kit's actual system prompt and a
realistic sample input to the model, prints the structured decision, and then
executes any proposed actions against MOCK tools (labelled [MOCK]).

The model/reasoning is real. The tools are stubs that return fake data and do
NOT touch any real system. Wiring the mock tools to your own integrations
(CRM, database, GitHub, etc.) is the final implementation step — see README.md.

Works with Anthropic OR OpenAI — it auto-detects whichever API key is set.

© 2026 AgentKits · www.agent-kits.com · Licensed CC BY 4.0 (attribution required)
Original blueprint: https://agent-kits.com/kit/vendor-risk-assessor

Setup (pick one):
    pip install anthropic   &&  export ANTHROPIC_API_KEY=sk-ant-...
    pip install openai      &&  export OPENAI_API_KEY=sk-...
Then:
    python run.py
"""
import os, json, sys

ANTHROPIC_MODEL = "claude-sonnet-4-6"
OPENAI_MODEL = "gpt-4o"

SYSTEM_PROMPT = """You are a Vendor Risk Assessment Agent supporting a third-party risk (TPRM) team. You produce a first-pass, evidence-based risk assessment of ONE vendor and recommend a risk tier or escalation. You assist the decision; you are NOT the final approver. You are judged on accuracy, evidence discipline, and never clearing a risky vendor or fabricating assurance.

== CORE PRINCIPLES ==
1. Evidence or it's a gap. Base every finding on a source you actually reviewed (questionnaire answer, SOC 2/ISO report, DPA, breach record) and cite it. If evidence for a control is missing, that is a GAP — never assume the control exists.
2. Risk is contextual. Weight risk by what the vendor will access: a vendor touching sensitive/customer data or critical systems is inherently higher risk than a low-touch tool, regardless of polish.
3. Recommend, don't approve. Output a tier and conditions for a human. High-risk and critical-data vendors are escalated, not cleared.

== HARD RULES (NON-NEGOTIABLE) ==
- NO AUTO-APPROVAL OF HIGH RISK: You may recommend fast-tracking only clearly LOW-risk vendors with no sensitive-data access and complete evidence. MODERATE/HIGH risk, critical-data access, or missing key evidence MUST be escalated to the risk team.
- NO FABRICATED ASSURANCE: Never claim a certification, control, or clean record you haven't seen evidence for. Expired or absent = flag it. Do not infer SOC 2 from a marketing page.
- MISSING EVIDENCE IS A FINDING: Treat absent DPA, expired SOC 2, unanswered security questions, or no breach-history check as explicit gaps with required remediation.
- NOT LEGAL/FINAL: You provide a risk view, not a legal opinion or contract sign-off.
- DATA HANDLING: Treat vendor and internal data as confidential; keep it in scope.

== ASSESSMENT DOMAINS ==
Security posture (SOC 2/ISO, pen-test, MFA/encryption); compliance (DPA, GDPR/regulatory fit, data residency); data access & scope (what data, how much, fourth parties/subprocessors); financial/operational stability; and incident/breach history. For each: state the evidence, the gap (if any), and a domain risk rating.

== DECISION POLICY ==
- LOW (fast-track recommend): no sensitive-data access, complete evidence, no material gaps. Recommend approval with standard terms (human confirms).
- MODERATE (conditions): real gaps that are remediable; list required conditions (e.g. obtain DPA, MFA enforcement) before/with onboarding.
- HIGH / ESCALATE: sensitive/critical-data access with gaps, prior unremediated breach, missing key evidence, or conflicting signals. Route to the risk team with findings.

== OUTPUT FORMAT (return ONE JSON object) ==
{
  "vendor": "<name>",
  "data_access": "<what they access; sensitive? critical?>",
  "overall_risk": "low|moderate|high",
  "recommendation": "fast_track|conditional|escalate",
  "domains": [ { "domain": "security|compliance|data_access|financial|breach_history", "evidence": "<source cited>", "gap": "<gap or 'none'>", "rating": "low|moderate|high" } ],
  "gaps": ["<material gaps / missing evidence>"],
  "conditions": ["<remediations required for onboarding, if conditional>"],
  "rationale": "<grounded summary>",
  "escalation": { "needed": <bool>, "to": "risk_team|none", "reason": "<why, or empty>" },
  "not_final_approval": true
}
If key evidence is missing or data access is sensitive/critical with any gap, set recommendation to escalate. Never claim assurance without cited evidence.
"""

SAMPLE_INPUT = """Vendor: MeetEasy (scheduling). Data access: calendar metadata only, no PII/customer data. Evidence: SOC 2 Type II (current, 2026), questionnaire complete, MFA + encryption attested, no breach history.
"""

# ---- MOCK tools (stubs — replace with real integrations to deploy) ----
def get_vendor_profile(args=None):
    print(f"  [MOCK] get_vendor_profile(" + str(args) + ")")
    return {"status": "mock", "tool": "get_vendor_profile"}

def questionnaire_parse(args=None):
    print(f"  [MOCK] questionnaire_parse(" + str(args) + ")")
    return {"status": "mock", "tool": "questionnaire_parse"}

def cert_verify(args=None):
    print(f"  [MOCK] cert_verify(" + str(args) + ")")
    return {"status": "mock", "tool": "cert_verify"}

def breach_history_lookup(args=None):
    print(f"  [MOCK] breach_history_lookup(" + str(args) + ")")
    return {"status": "mock", "tool": "breach_history_lookup"}

def data_scope_assess(args=None):
    print(f"  [MOCK] data_scope_assess(" + str(args) + ")")
    return {"status": "mock", "tool": "data_scope_assess"}

def score_risk_domains(args=None):
    print(f"  [MOCK] score_risk_domains(" + str(args) + ")")
    return {"status": "mock", "tool": "score_risk_domains"}

def recommend_tier(args=None):
    print(f"  [MOCK] recommend_tier(" + str(args) + ")")
    return {"status": "mock", "tool": "recommend_tier"}

def escalate_to_risk_team(args=None):
    print(f"  [MOCK] escalate_to_risk_team(" + str(args) + ")")
    return {"status": "mock", "tool": "escalate_to_risk_team"}

TOOLS = {
    "get_vendor_profile": get_vendor_profile,
    "questionnaire_parse": questionnaire_parse,
    "cert_verify": cert_verify,
    "breach_history_lookup": breach_history_lookup,
    "data_scope_assess": data_scope_assess,
    "score_risk_domains": score_risk_domains,
    "recommend_tier": recommend_tier,
    "escalate_to_risk_team": escalate_to_risk_team,
}


def call_model(system, user):
    """Dispatch to whichever provider's key is configured."""
    if os.environ.get("ANTHROPIC_API_KEY"):
        try:
            from anthropic import Anthropic
        except ImportError:
            sys.exit("Run: pip install anthropic")
        msg = Anthropic().messages.create(
            model=ANTHROPIC_MODEL,
            max_tokens=1600,
            system=system,
            messages=[{"role": "user", "content": user}],
        )
        return "".join(b.text for b in msg.content if getattr(b, "type", "") == "text")
    if os.environ.get("OPENAI_API_KEY"):
        try:
            from openai import OpenAI
        except ImportError:
            sys.exit("Run: pip install openai")
        resp = OpenAI().chat.completions.create(
            model=OPENAI_MODEL,
            max_tokens=1600,
            messages=[
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
        )
        return resp.choices[0].message.content
    sys.exit("Set ANTHROPIC_API_KEY or OPENAI_API_KEY first.")


def extract_json(text):
    a, b = text.find("{"), text.rfind("}")
    if a == -1 or b == -1:
        return None
    try:
        return json.loads(text[a : b + 1])
    except json.JSONDecodeError:
        return None


def main():
    print("=" * 70)
    print("INPUT:\n" + SAMPLE_INPUT)
    print("=" * 70)

    text = call_model(SYSTEM_PROMPT, SAMPLE_INPUT)
    print("AGENT DECISION:\n" + text)
    print("=" * 70)

    data = extract_json(text)
    if not data:
        print("(No JSON object parsed — the model returned prose.)")
        return

    actions = data.get("actions") or []
    if actions:
        print("EXECUTING PROPOSED ACTIONS (mock):")
        for a in actions:
            name = a.get("tool")
            args = a.get("args", {})
            if name in TOOLS:
                TOOLS[name](args)
            else:
                print(f"  [skip] no mock tool for '{name}'")
    else:
        print("This blueprint returned a structured result (no side-effecting "
              "actions in this sample). Available mock tools: "
              + ", ".join(TOOLS.keys()))


if __name__ == "__main__":
    main()
