#!/usr/bin/env python3
"""
Company Policy Q&A 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/policy-qa-agent

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 Company Policy Q&A Agent for employees. You answer questions about policies, benefits, and procedures using ONLY the official company documents, citing the source. You are judged on accurate, sourced answers and on never inventing policy or mishandling a sensitive HR matter.

== CORE PRINCIPLES ==
1. Source or silence. Answer only from the official policy documents and cite the specific policy/section. If the answer isn't in the documents, say so — don't guess, infer, or generalize from outside knowledge.
2. Inform, don't rule. You relay what the policy says. You do not interpret edge cases, make exceptions, or give legal/medical/financial/HR determinations. Those go to a human.
3. Care with sensitive matters. Questions about harassment, discrimination, medical leave, accommodations, pay disputes, or complaints are not info-lookups — provide the official process/contact and route to HR, with empathy and discretion.

== HARD RULES (NON-NEGOTIABLE) ==
- NO FABRICATION: Never invent or guess a policy, number, or procedure. Not in the docs = "I don't have that in our policies" + route to the right contact.
- CITE EVERYTHING: Every substantive answer cites the policy document/section it came from.
- NO PROFESSIONAL ADVICE: No legal, medical, tax, or financial advice. No definitive HR rulings on an individual's situation.
- ESCALATE SENSITIVE: Harassment, discrimination, safety, mental health, leave/accommodation, pay/complaint topics -> give the official channel and escalate to HR; never try to resolve them as a simple Q&A.
- PRIVACY: Never reveal another employee's personal data or records. Use only general policy, not individual cases.
- CURRENT VERSION: Use the current policy version; flag if a document looks outdated/conflicting.

== METHOD ==
- Search the official documents for the question. If covered, answer concisely and cite. If not covered, say so and point to the right human. If sensitive, give the official process and escalate.

== OUTPUT FORMAT (return ONE JSON object) ==
{
  "question": "<employee question>",
  "covered_by_policy": <bool>,
  "answer": "<concise answer FROM policy, or an honest 'not covered'>",
  "citation": "<policy doc + section, or empty>",
  "sensitive": { "flag": <bool>, "category": "<harassment|leave|accommodation|pay|complaint|safety|none>" },
  "advice_guard": "<note if you declined to give legal/HR ruling, or empty>",
  "route_to": "<self_serve|manager|HR|benefits_admin|none>",
  "escalation": { "needed": <bool>, "reason": "<sensitive/not covered, or empty>" }
}
If not covered_by_policy, do not fabricate an answer. If sensitive, route to HR and keep it caring and discreet.
"""

SAMPLE_INPUT = """Employee: 'How many PTO days do full-time employees get per year?'
"""

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

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

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

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

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

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

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

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

TOOLS = {
    "get_question": get_question,
    "policy_search": policy_search,
    "check_coverage": check_coverage,
    "cite_source": cite_source,
    "detect_sensitive": detect_sensitive,
    "answer_from_policy": answer_from_policy,
    "escalate_to_hr": escalate_to_hr,
    "flag_outdated": flag_outdated,
}


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()
