LangChain Agent Guardrails with Veto

from langchain_openai import ChatOpenAI
from langchain.agents import create_react_agent, AgentExecutor
from langchain_core.tools import tool

from veto import Veto, VetoOptions
from veto.integrations.langchain import VetoMiddleware

# Define your LangChain tools
@tool
def send_email(to: str, subject: str, body: str) -> str:
    """Send an email to a recipient."""
    # Your email sending logic
    return f"Email sent to {to}"

@tool
def process_payment(amount: float, recipient: str) -> str:
    """Process a payment transaction."""
    # Your payment processing logic
    return f"Payment of ${amount} processed"

# Initialize Veto
veto = await Veto.init(VetoOptions(
    api_key="veto_live_xxx",  # Or set VETO_API_KEY env var
))

# Create the middleware
middleware = VetoMiddleware(
    veto,
    on_deny=lambda name, args, reason: print(f"Blocked {name}: {reason}"),
)

# Create agent with guardrails
llm = ChatOpenAI(model="gpt-4o")
tools = [send_email, process_payment]

agent = create_react_agent(llm, tools)
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    middleware=[middleware],  # Add Veto middleware
)

# Run the agent - tool calls are now validated
result = await agent_executor.ainvoke({
    "input": "Send a $5000 payment to vendor@example.com"
})

from langgraph.graph import StateGraph, MessagesState
from langgraph.prebuilt import ToolNode

from veto import Veto, VetoOptions
from veto.integrations.langchain import create_veto_tool_node

# Initialize Veto
veto = await Veto.init(VetoOptions(api_key="veto_live_xxx"))

# Create your tools and ToolNode
tools = [send_email, process_payment, query_database]
tool_node = ToolNode(tools)

# Wrap with Veto for authorization
veto_tool_node = create_veto_tool_node(veto, tool_node)

# Build your LangGraph
workflow = StateGraph(MessagesState)
workflow.add_node("tools", veto_tool_node)
workflow.add_node("agent", agent_node)
# ... rest of your graph definition

from langchain.tools import tool
from langchain.agents.middleware import wrap_tool_call
from langchain.tools.tool_node import ToolCallRequest
from langchain.messages import ToolMessage
from typing import Callable

# Permission store (in production, use a database)
PERMISSIONS = {
    "admin": ["delete_data", "write_data", "read_data"],
    "editor": ["write_data", "read_data"],
    "viewer": ["read_data"]
}

@wrap_tool_call
def authorize_tool_call(
    request: ToolCallRequest,
    handler: Callable[[ToolCallRequest], ToolMessage],
) -> ToolMessage:
    """Authorization middleware that checks permissions before tool execution."""
    tool_name = request.tool_call["name"]
    user_permissions = request.config.get("configurable", {}).get("permissions", [])

    if tool_name not in user_permissions:
        return ToolMessage(
            content=f"Access denied: You don't have permission to use '{tool_name}'",
            name=tool_name,
            tool_call_id=request.tool_call["id"]
        )

    print(f"[Auth] Authorized tool call: {tool_name}")
    return handler(request)

# Using with an agent
from langchain.agents import create_agent

agent = create_agent(
    model="gpt-4",
    tools=[read_data, write_data, delete_data],
    middleware=[authorize_tool_call],
)

from dataclasses import dataclass
from langchain.agents import create_agent
from langchain.agents.middleware import wrap_model_call, ModelRequest, ModelResponse
from typing import Callable

@dataclass
class UserContext:
    user_id: str
    role: str
    permissions: list[str]

@wrap_model_call
def role_based_tool_filter(
    request: ModelRequest,
    handler: Callable[[ModelRequest], ModelResponse]
) -> ModelResponse:
    """Filter tools based on user role from runtime context."""
    context: UserContext = request.runtime.context
    role = context.role

    # Define role-based tool access
    role_tools = {
        "admin": request.tools,  # All tools
        "editor": [t for t in request.tools if t.name != "delete_data"],
        "viewer": [t for t in request.tools if t.name.startswith("read_")]
    }

    filtered_tools = role_tools.get(role, [])
    request = request.override(tools=filtered_tools)

    print(f"[RBAC] User {context.user_id} with role '{role}' has access to tools: {[t.name for t in filtered_tools]}")

    return handler(request)

# Create agent with context schema
agent = create_agent(
    model="gpt-4",
    tools=[read_data, write_data, delete_data],
    middleware=[role_based_tool_filter],
    context_schema=UserContext
)

# Invoke with user context
result = agent.invoke(
    {"messages": [{"role": "user", "content": "Delete all records"}]},
    context=UserContext(user_id="user-123", role="editor", permissions=["read", "write"])
)

from langchain.tools import tool
from langgraph.types import interrupt, Command
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import InMemorySaver

@tool
def send_email(to: str, subject: str, body: str) -> str:
    """Send an email to a recipient. Requires approval."""
    # Pause for human approval
    response = interrupt({
        "action": "send_email",
        "to": to,
        "subject": subject,
        "body": body,
        "message": "Approve sending this email?"
    })

    if response.get("action") == "approve":
        # Resume with potentially modified values
        final_to = response.get("to", to)
        final_subject = response.get("subject", subject)
        final_body = response.get("body", body)
        return f"Email sent to {final_to} with subject '{final_subject}'"
    return "Email cancelled by user"

@tool
def transfer_funds(from_account: str, to_account: str, amount: float) -> str:
    """Transfer funds between accounts. Requires approval."""
    response = interrupt({
        "action": "transfer_funds",
        "from": from_account,
        "to": to_account,
        "amount": amount,
        "risk_level": "high" if amount > 10000 else "medium"
    })

    if response.get("approved"):
        return f"Transferred ${response.get('amount', amount)} from {from_account} to {to_account}"
    return "Transfer cancelled"

# Create agent with checkpointer for state persistence
checkpointer = InMemorySaver()
agent = create_react_agent(
    model="openai:gpt-4",
    tools=[send_email, transfer_funds],
    checkpointer=checkpointer
)

# Invoke and handle interrupt
config = {"configurable": {"thread_id": "session-123"}}
result = agent.invoke(
    {"messages": [{"role": "user", "content": "Send email to john@example.com about the meeting"}]},
    config
)

# Check if there's an interrupt
if result.get("__interrupt__"):
    print(f"Approval needed: {result['__interrupt__']}")
    # Resume with approval
    result = agent.invoke(
        Command(resume={"action": "approve"}),
        config
    )

from langchain_core.tools import tool
from langchain_core.messages import ToolMessage, AIMessage
from langgraph.prebuilt import ToolNode
from langgraph.graph import StateGraph, MessagesState, START, END
from typing import Literal

@tool
def read_database(query: str) -> str:
    """Execute a read-only database query."""
    return f"Query result: {query}"

@tool
def write_database(table: str, data: dict) -> str:
    """Write data to database table."""
    return f"Wrote to {table}: {data}"

@tool
def delete_records(table: str, condition: str) -> str:
    """Delete records from table."""
    return f"Deleted from {table} where {condition}"

class AuthorizedToolNode:
    """ToolNode wrapper with built-in authorization."""

    def __init__(self, tools, permission_checker=None):
        self.tool_node = ToolNode(tools)
        self.tools_by_name = {t.name: t for t in tools}
        self.permission_checker = permission_checker or self.default_permission_check

    def default_permission_check(self, tool_name: str, user_context: dict) -> bool:
        """Default RBAC permission check."""
        role_permissions = {
            "admin": {"read_database", "write_database", "delete_records"},
            "editor": {"read_database", "write_database"},
            "viewer": {"read_database"}
        }
        role = user_context.get("role", "viewer")
        return tool_name in role_permissions.get(role, set())

    def __call__(self, state: MessagesState, config: dict = None) -> dict:
        """Execute tools with authorization check."""
        messages = state["messages"]
        last_message = messages[-1]

        if not isinstance(last_message, AIMessage) or not last_message.tool_calls:
            return {"messages": []}

        user_context = config.get("configurable", {}).get("user_context", {})
        tool_messages = []

        for tool_call in last_message.tool_calls:
            tool_name = tool_call["name"]

            # Authorization check
            if not self.permission_checker(tool_name, user_context):
                tool_messages.append(ToolMessage(
                    content=f"Authorization denied for tool '{tool_name}'. Contact administrator.",
                    tool_call_id=tool_call["id"],
                    name=tool_name
                ))
                continue

            # Execute authorized tool
            result = self.tool_node.invoke({"messages": [AIMessage(content="", tool_calls=[tool_call])]})
            tool_messages.extend(result.get("messages", []))

        return {"messages": tool_messages}

# Use in a graph
tools = [read_database, write_database, delete_records]
authorized_tool_node = AuthorizedToolNode(tools)

builder = StateGraph(MessagesState)
builder.add_node("tools", authorized_tool_node)
builder.add_edge(START, "tools")
builder.add_edge("tools", END)

graph = builder.compile()

# Invoke with user context
result = graph.invoke(
    {"messages": [AIMessage(content="", tool_calls=[
        {"id": "1", "name": "delete_records", "args": {"table": "users", "condition": "id=1"}}
    ])]},
    config={"configurable": {"user_context": {"role": "editor", "user_id": "user-123"}}}
)