"""
LangChain + LLM Shield (Production Architecture)

LiteLLM proxy handles input/output guardrails automatically.
App only calls Shield for tool RBAC + data policy enforcement.
"""
import os, json, requests
from langchain_openai import ChatOpenAI
from langchain.tools import tool
from langchain_core.messages import HumanMessage, ToolMessage

# LiteLLM proxy — your LLM gateway (guardrails run here automatically)
LITELLM_URL = os.getenv("LITELLM_URL", "https://litellm.your-company.com")
# Shield — tool RBAC + data policies (separate service)
SHIELD_URL  = os.getenv("LLM_SHIELD_URL", "https://shield.votal.ai")
API_KEY     = os.getenv("API_KEY")
AGENT_ID    = os.getenv("AGENT_ID", "my-langchain-agent")
USER_ROLE   = os.getenv("USER_ROLE", "user")

shield = requests.Session()
shield.headers.update({
    "X-API-Key":   API_KEY,
    "X-Agent-Key": AGENT_ID,
    "X-User-Role": USER_ROLE,
    "Content-Type": "application/json",
})


# ── 1. Register Agent (do this once) ──────────────────────────────
def register_agent():
    """Register the agent so Shield knows about it.
    If you skip this, the agent appears as a SHADOW AGENT."""
    resp = shield.post(f"{SHIELD_URL}/v1/agents/registry", json={
        "agent_id": AGENT_ID,
        "name": "LangChain Support Agent",
        "tools": ["search_faq", "create_ticket", "check_order_status"],
        "role_permissions": {
            "user":    ["search_faq", "check_order_status"],
            "support": ["search_faq", "create_ticket", "check_order_status"],
            "admin":   ["search_faq", "create_ticket", "check_order_status"],
        },
    })
    print("Registered:", resp.json())


# ── 2. Define LangChain tools ─────────────────────────────────────
@tool
def search_faq(query: str) -> str:
    """Search the FAQ knowledge base for answers."""
    return f"FAQ result for '{query}': Check our help center at /help."

@tool
def create_ticket(subject: str, description: str) -> str:
    """Create a support ticket for the customer."""
    return f"Ticket created: {subject} — ID: TKT-{hash(subject) % 10000}"

@tool
def check_order_status(order_id: str) -> str:
    """Check the current status of a customer order."""
    return f"Order {order_id}: Shipped, arriving in 2 days."

lc_tools = [search_faq, create_ticket, check_order_status]


# ── 3. Shielded LangChain agent ──────────────────────────────────
def run_agent(user_message: str):
    # LangChain LLM pointed at LiteLLM proxy
    # Input/output guardrails are applied automatically by LiteLLM
    llm = ChatOpenAI(
        model="default",
        base_url=f"{LITELLM_URL}/v1",
        api_key=os.getenv("OPENAI_API_KEY", "not-needed"),
        default_headers={
            "X-API-Key": API_KEY,
            "X-Agent-Key": AGENT_ID,
            "X-User-Role": USER_ROLE,
        },
    )
    llm_with_tools = llm.bind_tools(lc_tools)

    # Step 1: LLM call via LiteLLM (guardrails automatic)
    messages = [HumanMessage(content=user_message)]
    try:
        response = llm_with_tools.invoke(messages)
    except Exception as e:
        return f"Blocked by guardrails: {e}"

    if not response.tool_calls:
        return response.content

    # Step 2: Process tool calls with Shield RBAC + data policies
    messages.append(response)
    tool_map = {t.name: t for t in lc_tools}
    output_parts = []

    for tc in response.tool_calls:
        name, args = tc["name"], tc["args"]

        # 2a. Shield RBAC + data policy check on tool args
        check = shield.post(f"{SHIELD_URL}/v1/shield/tool/check", json={
            "tool_name": name,
            "tool_input": args,
            "agent_key": AGENT_ID,
            "user_role": USER_ROLE,
        }).json()
        if not check.get("allowed", True):
            reason = check.get("reason", "denied by RBAC")
            messages.append(ToolMessage(content=f"DENIED: {reason}",
                                        tool_call_id=tc["id"]))
            output_parts.append(f"BLOCKED: {name} — {reason}")
            continue

        # 2b. Execute the tool locally
        raw_output = tool_map[name].invoke(args)

        # 2c. Sanitize tool output (mask/redact/block per data policy)
        sanitized = shield.post(f"{SHIELD_URL}/v1/shield/tool/output", json={
            "tool_name": name,
            "tool_output": str(raw_output),
            "agent_key": AGENT_ID,
        }).json()
        if sanitized.get("action") == "block":
            messages.append(ToolMessage(
                content="Tool output blocked by data policy.",
                tool_call_id=tc["id"]))
        else:
            clean = sanitized.get("sanitized_output", str(raw_output))
            messages.append(ToolMessage(content=clean, tool_call_id=tc["id"]))
            output_parts.append(f"[{name}] {clean}")

    # Step 3: Final LLM call (output guardrails automatic via LiteLLM)
    try:
        final = llm_with_tools.invoke(messages)
        return final.content
    except Exception:
        return "\n".join(output_parts)


if __name__ == "__main__":
    # register_agent()  # Uncomment on first run
    print(run_agent("What is your return policy?"))
    print(run_agent("Check status of order ORD-555"))
    print(run_agent("Create a ticket: billing issue"))

"""
CrewAI + LLM Shield (Production Architecture)

LiteLLM proxy handles input/output guardrails automatically.
App only calls Shield for per-agent tool RBAC + data policies.
"""
import os, json, requests
from crewai import Agent, Task, Crew, Process
from crewai.tools import BaseTool
from typing import Type
from pydantic import BaseModel, Field

# LiteLLM proxy — your LLM gateway (guardrails run here automatically)
LITELLM_URL = os.getenv("LITELLM_URL", "https://litellm.your-company.com")
# Shield — tool RBAC + data policies (separate service)
SHIELD_URL  = os.getenv("LLM_SHIELD_URL", "https://shield.votal.ai")
API_KEY     = os.getenv("API_KEY")

shield = requests.Session()
shield.headers.update({
    "X-API-Key":   API_KEY,
    "Content-Type": "application/json",
})


# ── 1. Register all crew agents (do once) ─────────────────────────
def register_crew():
    agents = [
        {
            "agent_id": "research-agent",
            "name": "Research Analyst",
            "tools": ["web_search", "document_search"],
            "role_permissions": {
                "analyst": ["web_search", "document_search"],
                "viewer":  ["web_search"],
            },
        },
        {
            "agent_id": "writer-agent",
            "name": "Report Writer",
            "tools": ["generate_report", "send_email"],
            "role_permissions": {
                "analyst": ["generate_report", "send_email"],
                "viewer":  ["generate_report"],
            },
        },
    ]
    for a in agents:
        r = shield.post(f"{SHIELD_URL}/v1/agents/registry", json=a)
        print(f"Register {a['agent_id']}: {r.status_code}")


# ── 2. Shield-aware tool wrapper ──────────────────────────────────
class ShieldTool(BaseTool):
    """Wraps tool calls with Shield RBAC + data policy checks.
    Input/output guardrails are handled by LiteLLM automatically."""
    agent_key: str = ""
    user_role: str = "analyst"

    def _run(self, query: str) -> str:
        # PRE-EXECUTION: RBAC + data policy check on tool args
        tool_check = shield.post(f"{SHIELD_URL}/v1/shield/tool/check", json={
            "tool_name": self.name,
            "tool_input": {"query": query},
            "agent_key": self.agent_key,
            "user_role": self.user_role,
        }).json()
        if not tool_check.get("allowed", True):
            return f"BLOCKED: {self.name} — {tool_check.get('reason', 'denied')}"

        # Execute tool (your actual implementation here)
        raw_output = f"[{self.name}] Results for: {query}"

        # POST-EXECUTION: Sanitize tool output (mask/redact/block)
        sanitized = shield.post(f"{SHIELD_URL}/v1/shield/tool/output", json={
            "tool_name": self.name,
            "tool_output": raw_output,
            "agent_key": self.agent_key,
        }).json()
        if sanitized.get("action") == "block":
            return f"BLOCKED: {self.name} output violated data policy"
        return sanitized.get("sanitized_output", raw_output)


# ── 3. Build the crew ─────────────────────────────────────────────
class WebSearchInput(BaseModel):
    query: str = Field(description="Search query")

class WebSearchTool(ShieldTool):
    name: str = "web_search"
    description: str = "Search the web for information"
    args_schema: Type[BaseModel] = WebSearchInput
    agent_key: str = "research-agent"

class GenerateReportInput(BaseModel):
    query: str = Field(description="Report topic and data")

class GenerateReportTool(ShieldTool):
    name: str = "generate_report"
    description: str = "Generate a formatted report"
    args_schema: Type[BaseModel] = GenerateReportInput
    agent_key: str = "writer-agent"

def run_crew(topic: str, user_role: str = "analyst"):
    # CrewAI uses LiteLLM as its LLM backend
    # Input/output guardrails are applied automatically by LiteLLM
    researcher = Agent(
        role="Research Analyst",
        goal=f"Research: {topic}",
        backstory="Senior analyst with deep domain expertise",
        tools=[WebSearchTool(user_role=user_role)],
        llm=f"litellm/{LITELLM_URL}/v1/default",
        verbose=True,
    )
    writer = Agent(
        role="Report Writer",
        goal="Write a clear, actionable report",
        backstory="Technical writer specializing in research reports",
        tools=[GenerateReportTool(user_role=user_role)],
        llm=f"litellm/{LITELLM_URL}/v1/default",
        verbose=True,
    )

    crew = Crew(
        agents=[researcher, writer],
        tasks=[
            Task(description=f"Research {topic}", expected_output="Findings", agent=researcher),
            Task(description="Write report", expected_output="Report", agent=writer),
        ],
        process=Process.sequential,
        verbose=True,
    )
    return str(crew.kickoff())


if __name__ == "__main__":
    # register_crew()  # Uncomment on first run
    print(run_crew("Analyze Q1 2026 market trends in AI"))

"""
OpenAI Agents SDK + LLM Shield (Production Architecture)

LiteLLM proxy handles input/output guardrails automatically.
App only calls Shield for tool RBAC + data policy enforcement.
"""
import os, json, requests
from openai import OpenAI

# LiteLLM proxy — your LLM gateway (guardrails run here automatically)
LITELLM_URL = os.getenv("LITELLM_URL", "https://litellm.your-company.com")
# Shield — tool RBAC + data policies (separate service)
SHIELD_URL  = os.getenv("LLM_SHIELD_URL", "https://shield.votal.ai")
API_KEY     = os.getenv("API_KEY")
AGENT_ID    = os.getenv("AGENT_ID", "my-openai-agent")
USER_ROLE   = os.getenv("USER_ROLE", "user")

# OpenAI client pointed at LiteLLM (guardrails automatic)
client = OpenAI(base_url=f"{LITELLM_URL}/v1",
                api_key=os.getenv("OPENAI_API_KEY", "not-needed"),
                default_headers={"X-API-Key": API_KEY,
                                 "X-Agent-Key": AGENT_ID,
                                 "X-User-Role": USER_ROLE})
shield = requests.Session()
shield.headers.update({
    "X-API-Key":   API_KEY,
    "X-Agent-Key": AGENT_ID,
    "X-User-Role": USER_ROLE,
    "Content-Type": "application/json",
})


# ── 1. Register Agent (do this once) ──────────────────────────────
def register_agent():
    shield.post(f"{SHIELD_URL}/v1/agents/registry", json={
        "agent_id": AGENT_ID,
        "name": "My OpenAI Agent",
        "tools": ["lookup_order", "cancel_order", "get_refund_status"],
        "role_permissions": {
            "user":  ["lookup_order", "get_refund_status"],
            "admin": ["lookup_order", "cancel_order", "get_refund_status"],
        },
    })


# ── 2. Define tools (OpenAI function calling format) ──────────────
tools = [
    {"type": "function", "function": {
        "name": "lookup_order",
        "description": "Look up order details by order ID",
        "strict": True,
        "parameters": {
            "type": "object",
            "properties": {"order_id": {"type": "string"}},
            "required": ["order_id"],
            "additionalProperties": False,
        },
    }},
    {"type": "function", "function": {
        "name": "cancel_order",
        "description": "Cancel an existing order",
        "strict": True,
        "parameters": {
            "type": "object",
            "properties": {"order_id": {"type": "string"}, "reason": {"type": "string"}},
            "required": ["order_id", "reason"],
            "additionalProperties": False,
        },
    }},
]


# ── 3. Shielded agent loop ────────────────────────────────────────
def run_agent(user_message: str):
    # Step 1: LLM call via LiteLLM (input guardrails automatic)
    messages = [
        {"role": "system", "content": "You are a helpful support agent."},
        {"role": "user", "content": user_message},
    ]
    response = client.chat.completions.create(
        model="default", messages=messages,
        tools=tools, tool_choice="auto",
    )
    msg = response.choices[0].message

    if not msg.tool_calls:
        return msg.content  # Output guardrails already applied by LiteLLM

    # Step 2: Process tool calls with Shield RBAC + data policies
    for tc in msg.tool_calls:
        name = tc.function.name
        args = json.loads(tc.function.arguments)

        # 2a. Shield RBAC + data policy check on tool args
        check = shield.post(f"{SHIELD_URL}/v1/shield/tool/check", json={
            "tool_name": name,
            "tool_input": args,
            "agent_key": AGENT_ID,
            "user_role": USER_ROLE,
        }).json()
        if not check.get("allowed", True):
            messages.append({"role": "tool", "tool_call_id": tc.id,
                             "content": f"DENIED: {check.get('reason', 'RBAC')}"})
            continue

        # 2b. Execute tool (your implementation)
        raw_output = execute_tool(name, args)

        # 2c. Sanitize output (mask/redact/block per data policy)
        sanitized = shield.post(f"{SHIELD_URL}/v1/shield/tool/output", json={
            "tool_name": name,
            "tool_output": str(raw_output),
            "agent_key": AGENT_ID,
        }).json()
        if sanitized.get("action") == "block":
            messages.append({"role": "tool", "tool_call_id": tc.id,
                             "content": "Output blocked by data policy."})
        else:
            clean = sanitized.get("sanitized_output", str(raw_output))
            messages.append({"role": "tool", "tool_call_id": tc.id,
                             "content": clean})

    # Step 3: Final LLM call (output guardrails automatic via LiteLLM)
    final = client.chat.completions.create(model="default", messages=messages)
    return final.choices[0].message.content


if __name__ == "__main__":
    # register_agent()  # Uncomment on first run
    print(run_agent("What's the status of order ORD-12345?"))
    print(run_agent("Cancel order ORD-12345, I changed my mind"))

"""
Anthropic Claude + LLM Shield (Production Architecture)

LiteLLM proxy handles input/output guardrails automatically.
App only calls Shield for tool RBAC + data policy enforcement.
"""
import os, json, requests
import anthropic

# LiteLLM proxy — your LLM gateway (guardrails run here automatically)
LITELLM_URL = os.getenv("LITELLM_URL", "https://litellm.your-company.com")
# Shield — tool RBAC + data policies (separate service)
SHIELD_URL  = os.getenv("LLM_SHIELD_URL", "https://shield.votal.ai")
API_KEY     = os.getenv("API_KEY")
AGENT_ID    = os.getenv("AGENT_ID", "my-claude-agent")
USER_ROLE   = os.getenv("USER_ROLE", "user")

# Claude client via LiteLLM (guardrails automatic)
claude = anthropic.Anthropic(base_url=f"{LITELLM_URL}/v1",
                             api_key=os.getenv("ANTHROPIC_API_KEY", "not-needed"))
shield = requests.Session()
shield.headers.update({
    "X-API-Key":   API_KEY,
    "X-Agent-Key": AGENT_ID,
    "X-User-Role": USER_ROLE,
    "Content-Type": "application/json",
})


# ── 1. Register Agent ─────────────────────────────────────────────
def register_agent():
    shield.post(f"{SHIELD_URL}/v1/agents/registry", json={
        "agent_id": AGENT_ID,
        "name": "Claude Support Agent",
        "tools": ["search_knowledge_base", "create_ticket", "escalate_to_human"],
        "role_permissions": {
            "user":    ["search_knowledge_base"],
            "support": ["search_knowledge_base", "create_ticket"],
            "admin":   ["search_knowledge_base", "create_ticket", "escalate_to_human"],
        },
    })


# ── 2. Define tools (Anthropic format) ────────────────────────────
claude_tools = [
    {"name": "search_knowledge_base",
     "description": "Search the support knowledge base",
     "input_schema": {
         "type": "object",
         "properties": {"query": {"type": "string"}},
         "required": ["query"],
     }},
    {"name": "create_ticket",
     "description": "Create a support ticket",
     "input_schema": {
         "type": "object",
         "properties": {
             "subject": {"type": "string"},
             "description": {"type": "string"},
             "priority": {"type": "string", "enum": ["low", "medium", "high"]},
         },
         "required": ["subject", "description", "priority"],
     }},
]


# ── 3. Shield-wrapped Claude agent ────────────────────────────────
def run_claude_agent(user_message: str):
    # Step 1: Call Claude via LiteLLM (input guardrails automatic)
    response = claude.messages.create(
        model="claude-sonnet-4-20250514",
        max_tokens=1024,
        system="You are a helpful support agent. Use tools when needed.",
        messages=[{"role": "user", "content": user_message}],
        tools=claude_tools,
    )

    # Step 2: Process response blocks
    results = []
    tool_results = []
    for block in response.content:
        if block.type == "text":
            results.append(block.text)

        elif block.type == "tool_use":
            name = block.name

            # 2a. Shield RBAC + data policy check on tool args
            check = shield.post(f"{SHIELD_URL}/v1/shield/tool/check", json={
                "tool_name": name,
                "tool_input": block.input,
                "agent_key": AGENT_ID,
                "user_role": USER_ROLE,
            }).json()
            if not check.get("allowed", True):
                tool_results.append({"type": "tool_result", "tool_use_id": block.id,
                    "content": f"DENIED: {check.get('reason', 'RBAC')}"})
                continue

            # 2b. Execute tool (your implementation)
            raw_output = execute_tool(name, block.input)

            # 2c. Sanitize output (mask/redact/block per data policy)
            sanitized = shield.post(f"{SHIELD_URL}/v1/shield/tool/output", json={
                "tool_name": name,
                "tool_output": str(raw_output),
                "agent_key": AGENT_ID,
            }).json()
            if sanitized.get("action") == "block":
                tool_results.append({"type": "tool_result", "tool_use_id": block.id,
                    "content": "Output blocked by data policy."})
            else:
                clean = sanitized.get("sanitized_output", str(raw_output))
                tool_results.append({"type": "tool_result", "tool_use_id": block.id,
                    "content": clean})

    # Step 3: If tools were called, send results back to Claude
    if tool_results:
        final = claude.messages.create(
            model="claude-sonnet-4-20250514", max_tokens=1024,
            messages=[
                {"role": "user", "content": user_message},
                {"role": "assistant", "content": response.content},
                {"role": "user", "content": tool_results},
            ],
            tools=claude_tools,
        )
        return final.content[0].text

    return "\n".join(results)


if __name__ == "__main__":
    # register_agent()  # Uncomment on first run
    print(run_claude_agent("Search for how to reset my password"))
    print(run_claude_agent("Create an urgent ticket: billing error"))