From ce076216d93d0f50de645756aa9ed91b270d9640 Mon Sep 17 00:00:00 2001 From: Riyan Dhiman Date: Sat, 28 Mar 2026 03:26:37 +0530 Subject: [PATCH] examples --- examples/.claude/settings.json | 16 +++ examples/anthropic_example.py | 108 ++++++++++++++++ examples/claude_code_setup.py | 84 ++++++++----- examples/guard_decorator_example.py | 186 ++++++++++++++++++++-------- examples/langchain_example.py | 165 ++++++++++++++---------- examples/observe_mode_example.py | 102 +++++++++++++++ examples/openai_example.py | 119 ++++++++++++++++++ examples/openrouter_example.py | 87 +++++++++++++ examples/quickstart.py | 103 +++++++++------ 9 files changed, 784 insertions(+), 186 deletions(-) create mode 100644 examples/.claude/settings.json create mode 100644 examples/anthropic_example.py create mode 100644 examples/observe_mode_example.py create mode 100644 examples/openai_example.py create mode 100644 examples/openrouter_example.py diff --git a/examples/.claude/settings.json b/examples/.claude/settings.json new file mode 100644 index 0000000..830cad2 --- /dev/null +++ b/examples/.claude/settings.json @@ -0,0 +1,16 @@ +{ + "hooks": { + "PreToolUse": [ + { + "matcher": "", + "hooks": [ + { + "type": "command", + "command": "/Users/riyandhiman/Library/Python/3.9/bin/agsec check --format=claude-code --policy-dir /Users/riyandhiman/project/agsec/examples/policies", + "timeout": 30 + } + ] + } + ] + } +} \ No newline at end of file diff --git a/examples/anthropic_example.py b/examples/anthropic_example.py new file mode 100644 index 0000000..16e5aa5 --- /dev/null +++ b/examples/anthropic_example.py @@ -0,0 +1,108 @@ +""" +Real Anthropic example with agsec protection. + +Run: pip install agsec[anthropic] && export ANTHROPIC_API_KEY=sk-ant-... + python examples/anthropic_example.py +""" + +import json + +from anthropic import Anthropic + +from agsec.integrations.anthropic import protect, allow, deny, review, param + +# Protect the client — one line +client = protect( + Anthropic(), + allow("get_weather", "search"), + deny("delete_user"), + deny("run_sql").when(param("query").contains("DROP")), + review("send_email"), +) + +# Define tools +tools = [ + { + "name": "get_weather", + "description": "Get weather for a city", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + { + "name": "delete_user", + "description": "Delete a user account", + "input_schema": { + "type": "object", + "properties": {"user_id": {"type": "string"}}, + "required": ["user_id"], + }, + }, + { + "name": "run_sql", + "description": "Run a SQL query", + "input_schema": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + }, +] + +print("=== Test 1: Safe request (weather) ===") +response = client.messages.create( + model="claude-sonnet-4-20250514", + max_tokens=1024, + messages=[{"role": "user", "content": "What's the weather in Paris?"}], + tools=tools, +) + +for block in response.content: + if block.type == "tool_use": + print(f" ALLOWED: {block.name}({json.dumps(block.input)})") + elif block.type == "text": + print(f" Text: {block.text}") + +if hasattr(response, "_agsec_blocked") and response._agsec_blocked: + for b in response._agsec_blocked: + print(f" BLOCKED: {b['name']} — {b['reason']}") + + +print("\n=== Test 2: Dangerous request (delete user) ===") +response = client.messages.create( + model="claude-sonnet-4-20250514", + max_tokens=1024, + messages=[{"role": "user", "content": "Delete user xyz789 from the system"}], + tools=tools, +) + +for block in response.content: + if block.type == "tool_use": + print(f" ALLOWED: {block.name}({json.dumps(block.input)})") + elif block.type == "text": + print(f" Text: {block.text}") + +if hasattr(response, "_agsec_blocked") and response._agsec_blocked: + for b in response._agsec_blocked: + print(f" BLOCKED: {b['name']} — {b['reason']}") + + +print("\n=== Test 3: Conditional block (SQL injection) ===") +response = client.messages.create( + model="claude-sonnet-4-20250514", + max_tokens=1024, + messages=[{"role": "user", "content": "Run this query: DROP TABLE users"}], + tools=tools, +) + +for block in response.content: + if block.type == "tool_use": + print(f" ALLOWED: {block.name}({json.dumps(block.input)})") + elif block.type == "text": + print(f" Text: {block.text}") + +if hasattr(response, "_agsec_blocked") and response._agsec_blocked: + for b in response._agsec_blocked: + print(f" BLOCKED: {b['name']} — {b['reason']}") diff --git a/examples/claude_code_setup.py b/examples/claude_code_setup.py index 55d4258..676400b 100644 --- a/examples/claude_code_setup.py +++ b/examples/claude_code_setup.py @@ -1,42 +1,58 @@ """ -agsec + Claude Code — setup firewall in 30 seconds. +Claude Code + agsec setup — run this script to set everything up. -Run this script or use the CLI commands below. +Run: pip install agsec + python examples/claude_code_setup.py """ +import os import subprocess import sys -print(""" -=== agsec Claude Code Setup === -Three commands to protect your agent: - - 1. pip install agsec - 2. agsec init - 3. agsec install claude-code - -That's it. The firewall is active. - ---- What gets blocked by default --- - - rm -rf / BLOCKED (file deletion) - cat .env BLOCKED (secret access) - git push --force BLOCKED (force push) - git push origin main BLOCKED (protected branch) - curl --data secrets.json BLOCKED (data exfiltration) - ---- What gets allowed --- - - ls, grep, find ALLOWED (read ops) - python -m pytest ALLOWED (safe bash) - git push origin feature/x ALLOWED (feature branch) - ---- Manage policies --- - - agsec policy list # see all rules - agsec policy add # add a rule (interactive) - agsec policy remove # remove a rule - agsec validate # check for errors - agsec audit --stats # view activity -""") +def main(): + print("=== agsec + Claude Code Setup ===\n") + + # Step 1: Check if policies exist + has_policies = os.path.isdir("policies") or os.path.isdir(".agsec/policies") + + if not has_policies: + print("1. Creating policies...") + result = subprocess.run( + [sys.executable, "-m", "agsec", "init"], + capture_output=True, text=True, + ) + print(f" {result.stdout.strip().split(chr(10))[0]}") + else: + print("1. Policies already exist.") + + # Step 2: Install hook + print("\n2. Installing Claude Code hook...") + result = subprocess.run( + [sys.executable, "-m", "agsec", "install", "claude-code"], + capture_output=True, text=True, + ) + print(f" {result.stdout.strip().split(chr(10))[0]}") + + # Step 3: Show current policies + print("\n3. Active policies:") + result = subprocess.run( + [sys.executable, "-m", "agsec", "policy", "list"], + capture_output=True, text=True, + ) + for line in result.stdout.strip().split("\n"): + print(f" {line}") + + # Step 4: Validate + print("\n4. Validating...") + result = subprocess.run( + [sys.executable, "-m", "agsec", "validate"], + capture_output=True, text=True, + ) + print(f" {result.stdout.strip()}") + + print("\n=== Done. Restart Claude Code to activate the firewall. ===") + + +if __name__ == "__main__": + main() diff --git a/examples/guard_decorator_example.py b/examples/guard_decorator_example.py index 6ddbb5f..b3c259b 100644 --- a/examples/guard_decorator_example.py +++ b/examples/guard_decorator_example.py @@ -1,61 +1,145 @@ """ -agsec @guard decorator — protect any Python function with policies. +@guard decorator example — protect any Python function. No framework needed. Works with any code. + +Run: pip install agsec + python examples/guard_decorator_example.py """ +import asyncio import os -from agsec import guard - -# Point to policies directory -policy_dir = os.path.join(os.path.dirname(__file__), "policies") - - -# --- Protect functions with one decorator --- - -@guard("email.send", policy_dir=policy_dir) -def send_email(to, subject, body): - return {"sent_to": to, "subject": subject} - - -@guard("payment.charge", policy_dir=policy_dir) -def charge_payment(amount, recipient): - return {"charged": amount, "to": recipient} - - -@guard("db.read", policy_dir=policy_dir) -def read_database(table, query): - return {"table": table, "rows": 42} - - -# --- Async works too --- - -@guard("notification.push", policy_dir=policy_dir) -async def send_push(user_id, message): - return {"sent_to": user_id} +import tempfile +from agsec import guard +from agsec.exceptions import PolicyViolationError + + +def main(): + # Create a temp policy directory for this demo + with tempfile.TemporaryDirectory() as policy_dir: + # Write a simple policy + with open(os.path.join(policy_dir, "policy.yaml"), "w") as f: + f.write(""" +version: "1.0" +default: deny + +statements: + - sid: "BlockExternalEmail" + effect: deny + actions: ["email.send"] + conditions: + params.to: + op: "regex" + value: '.*@(?!company\.com$)' + reason: "External emails are blocked" + + - sid: "AllowEmail" + effect: allow + actions: ["email.send"] + + - sid: "AllowSmallPayments" + effect: allow + actions: ["payment.charge"] + conditions: + params.amount: + op: "<=" + value: 1000 + + - sid: "ReviewLargePayments" + effect: review + actions: ["payment.charge"] + conditions: + params.amount: + op: ">" + value: 1000 + reason: "Large payments need approval" + + - sid: "AllowRead" + effect: allow + actions: ["db.read"] + + - sid: "BlockWrite" + effect: deny + actions: ["db.write"] + reason: "Database writes are blocked" +""") + + # --- Decorate your functions --- + + @guard("email.send", policy_dir=policy_dir) + def send_email(to, subject, body): + return f"Sent to {to}: {subject}" + + @guard("payment.charge", policy_dir=policy_dir) + def charge_payment(amount, currency="USD"): + return f"Charged {currency} {amount}" + + @guard("db.read", policy_dir=policy_dir) + def read_db(table, query): + return f"Read {table}: {query}" + + @guard("db.write", policy_dir=policy_dir) + def write_db(table, data): + return f"Wrote to {table}" + + @guard("notification.push", policy_dir=policy_dir) + async def push_notification(user_id, message): + return f"Notified {user_id}: {message}" + + # --- Test them --- + + print("=== @guard Decorator Examples ===\n") + + # Internal email: ALLOWED + try: + result = send_email(to="alice@company.com", subject="Meeting", body="Hi!") + print(f" Internal email: ALLOWED -> {result}") + except PolicyViolationError as e: + print(f" Internal email: BLOCKED -> {e}") + + # External email: BLOCKED + try: + result = send_email(to="hacker@evil.com", subject="Secrets", body="...") + print(f" External email: ALLOWED -> {result}") + except PolicyViolationError as e: + print(f" External email: BLOCKED") + + # Small payment: ALLOWED + try: + result = charge_payment(amount=50) + print(f" Small payment: ALLOWED -> {result}") + except PolicyViolationError as e: + print(f" Small payment: BLOCKED") + + # Large payment: REVIEW (blocked with review reason) + try: + result = charge_payment(amount=50000) + print(f" Large payment: ALLOWED -> {result}") + except PolicyViolationError as e: + print(f" Large payment: REVIEW REQUIRED") + + # DB read: ALLOWED + try: + result = read_db(table="users", query="SELECT *") + print(f" DB read: ALLOWED -> {result}") + except PolicyViolationError as e: + print(f" DB read: BLOCKED") + + # DB write: BLOCKED + try: + result = write_db(table="users", data={"name": "hacker"}) + print(f" DB write: ALLOWED -> {result}") + except PolicyViolationError as e: + print(f" DB write: BLOCKED") + + # Async function: BLOCKED (no allow policy for notification.push) + try: + result = asyncio.run(push_notification(user_id="u1", message="hello")) + print(f" Push notification: ALLOWED -> {result}") + except PolicyViolationError as e: + print(f" Push notification: BLOCKED (default deny)") -# --- Try it out --- if __name__ == "__main__": - from agsec.exceptions import PolicyViolationError - - # Allowed - result = read_database(table="users", query="SELECT *") - print(f"DB read: {result}") - - # Allowed (small payment) - result = charge_payment(amount=50, recipient={"country": "US"}) - print(f"Small payment: {result}") - - # Blocked (large payment triggers review) - try: - charge_payment(amount=50000, recipient={"country": "US"}) - except PolicyViolationError as e: - print(f"Large payment: BLOCKED — {e}") - - # Blocked (email to external domain) - try: - send_email(to="hacker@evil.com", subject="secrets", body="...") - except PolicyViolationError as e: - print(f"External email: BLOCKED — {e}") + main() diff --git a/examples/langchain_example.py b/examples/langchain_example.py index a35c583..49f5494 100644 --- a/examples/langchain_example.py +++ b/examples/langchain_example.py @@ -1,70 +1,107 @@ """ -agsec + LangChain — protect agent tools with inline policies. +Real LangChain example with agsec protection. -Install: pip install agsec[langchain] langchain-core +Run: pip install agsec[langchain] langchain-openai + export OPENAI_API_KEY=sk-... + python examples/langchain_example.py """ from agsec.integrations.langchain import guard, allow, deny, review, param -# --- Assume these are your LangChain tools --- -# from langchain_core.tools import tool -# -# @tool -# def search(query: str) -> str: -# """Search the web.""" -# return f"Results for: {query}" -# -# @tool -# def send_email(to: str, subject: str, body: str) -> str: -# """Send an email.""" -# return f"Email sent to {to}" -# -# @tool -# def payment(amount: float, recipient: str) -> str: -# """Process a payment.""" -# return f"Charged {amount} to {recipient}" -# -# @tool -# def delete_record(table: str, id: int) -> str: -# """Delete a database record.""" -# return f"Deleted {id} from {table}" - -# --- Protect your tools with one line --- - -# protected_tools = guard( -# allow(search), # search is always fine -# review(send_email), # emails need human approval -# deny(delete_record), # never let the agent delete -# deny(payment).when(param("amount") > 10000), # block large payments -# ) - -# --- Pass to your agent --- - -# from langgraph.prebuilt import create_react_agent -# agent = create_react_agent(llm, protected_tools) - -# --- What happens --- -# agent calls search("python docs") -> ALLOWED -# agent calls send_email(...) -> BLOCKED (review required) -# agent calls delete_record(...) -> BLOCKED -# agent calls payment(amount=500) -> ALLOWED -# agent calls payment(amount=50000) -> BLOCKED (amount > 10000) - -print(""" -agsec LangChain integration example. - -Usage: - from agsec.integrations.langchain import guard, allow, deny, review, param - - protected_tools = guard( - allow(search, calculator), - review(send_email), - deny(delete_record), - deny(payment).when(param("amount") > 10000), - ) - - agent = create_react_agent(llm, protected_tools) - -Install langchain-core to run this example: - pip install agsec[langchain] langchain-core -""") +# --- Step 1: Define your tools --- + +try: + from langchain_core.tools import tool + + @tool + def search(query: str) -> str: + """Search the web for information.""" + return f"Results for: {query}" + + @tool + def calculator(expression: str) -> str: + """Evaluate a math expression.""" + try: + return str(eval(expression)) # noqa: S307 + except Exception as e: + return f"Error: {e}" + + @tool + def send_email(to: str, subject: str, body: str) -> str: + """Send an email to someone.""" + return f"Email sent to {to}: {subject}" + + @tool + def delete_database(table: str) -> str: + """Delete a database table.""" + return f"Deleted table: {table}" + + @tool + def process_payment(amount: float, recipient: str) -> str: + """Process a payment.""" + return f"Paid ${amount} to {recipient}" + + # --- Step 2: Protect tools with one line --- + + protected_tools = guard( + allow(search, calculator), + review(send_email), + deny(delete_database), + deny(process_payment).when(param("amount") > 10000), + ) + + print("=== agsec + LangChain ===\n") + print(f"Protected {len(protected_tools)} tools:\n") + for t in protected_tools: + print(f" {t.name}: {t.description}") + + # --- Step 3: Test each tool --- + + print("\n--- Testing tools ---\n") + + # search: ALLOWED + try: + result = search.invoke({"query": "python docs"}) + print(f" search('python docs') -> {result}") + except Exception as e: + print(f" search -> BLOCKED: {e}") + + # delete_database: BLOCKED by agsec + print() + from agsec.exceptions import PolicyViolationError + + try: + # Use the protected version + result = protected_tools[3].invoke({"table": "users"}) + print(f" delete_database('users') -> {result}") + except PolicyViolationError as e: + print(f" delete_database('users') -> BLOCKED: {e}") + + # process_payment small: ALLOWED + try: + result = protected_tools[4].invoke({"amount": 50.0, "recipient": "Bob"}) + print(f" process_payment($50) -> {result}") + except PolicyViolationError as e: + print(f" process_payment($50) -> BLOCKED: {e}") + + # process_payment large: BLOCKED + try: + result = protected_tools[4].invoke({"amount": 50000.0, "recipient": "Bob"}) + print(f" process_payment($50000) -> {result}") + except PolicyViolationError as e: + print(f" process_payment($50000) -> BLOCKED: {e}") + + # --- Step 4: Use with an agent (uncomment if you have langchain-openai) --- + + # from langchain_openai import ChatOpenAI + # from langgraph.prebuilt import create_react_agent + # + # llm = ChatOpenAI(model="gpt-4o-mini") + # agent = create_react_agent(llm, protected_tools) + # + # result = agent.invoke({"messages": [{"role": "user", "content": "What's 42 * 17?"}]}) + # print(result["messages"][-1].content) + +except ImportError: + print("Install langchain-core to run this example:") + print(" pip install agsec[langchain]") diff --git a/examples/observe_mode_example.py b/examples/observe_mode_example.py new file mode 100644 index 0000000..2376a8b --- /dev/null +++ b/examples/observe_mode_example.py @@ -0,0 +1,102 @@ +""" +Observe mode example — audit everything, block nothing. + +This shows the full observe → enforce workflow: +1. Start in observe mode +2. Agent runs normally, all actions logged +3. Check what would be blocked +4. Switch to enforce when confident + +Run: pip install agsec + python examples/observe_mode_example.py +""" + +import json +import os +import subprocess +import sys +import tempfile + +def run_agsec(*args, stdin_data=None): + """Run agsec CLI command and return output.""" + cmd = [sys.executable, "-m", "agsec.cli.main"] + list(args) + result = subprocess.run( + cmd, + input=stdin_data, + capture_output=True, + text=True, + ) + return result.stdout.strip(), result.stderr.strip(), result.returncode + + +def main(): + # Use a temp directory for this demo + original_dir = os.getcwd() + with tempfile.TemporaryDirectory() as tmpdir: + os.chdir(tmpdir) + + print("=== Observe Mode Demo ===\n") + + # Step 1: Init in observe mode + print("1. Initialize in observe mode:") + stdout, _, _ = run_agsec("init", "--observe") + print(f" {stdout.split(chr(10))[0]}") + print() + + # Step 2: Simulate agent actions + print("2. Simulate agent actions (all will be ALLOWED in observe mode):\n") + + actions = [ + ("Bash", {"command": "ls -la"}, "list files"), + ("Bash", {"command": "cat .env"}, "read secrets"), + ("Write", {"file_path": ".env", "content": "SECRET=x"}, "write to .env"), + ("Bash", {"command": "git push --force origin main"}, "force push"), + ("Read", {"file_path": "app.py"}, "read source"), + ] + + policy_dir = os.path.join(tmpdir, "policies") + for tool_name, tool_input, desc in actions: + data = json.dumps({ + "tool_name": tool_name, + "tool_input": tool_input, + }) + _, stderr, code = run_agsec( + "check", "--policy-dir", policy_dir, + stdin_data=data, + ) + status = "ALLOWED" if code == 0 else "BLOCKED" + print(f" {status} {desc:20s} ({tool_name}: {json.dumps(tool_input)[:50]})") + + # Step 3: Check audit stats + print("\n3. Check what would have been blocked:\n") + stdout, _, _ = run_agsec("audit", "--stats") + for line in stdout.split("\n"): + print(f" {line}") + + # Step 4: Switch to enforce + print("\n4. Switch to enforce mode:") + stdout, _, _ = run_agsec("enforce") + print(f" {stdout.split(chr(10))[0]}") + + # Step 5: Try the same actions again + print("\n5. Same actions, now in enforce mode:\n") + + for tool_name, tool_input, desc in actions: + data = json.dumps({ + "tool_name": tool_name, + "tool_input": tool_input, + }) + _, stderr, code = run_agsec( + "check", "--policy-dir", policy_dir, + stdin_data=data, + ) + status = "ALLOWED" if code == 0 else "BLOCKED" + print(f" {status} {desc:20s} ({tool_name}: {json.dumps(tool_input)[:50]})") + + os.chdir(original_dir) + + print("\n=== Done ===") + + +if __name__ == "__main__": + main() diff --git a/examples/openai_example.py b/examples/openai_example.py new file mode 100644 index 0000000..87d5449 --- /dev/null +++ b/examples/openai_example.py @@ -0,0 +1,119 @@ +""" +Real OpenAI example with agsec protection. + +Run: pip install agsec[openai] && export OPENAI_API_KEY=sk-... + python examples/openai_example.py +""" + +import json + +from openai import OpenAI + +from agsec.integrations.openai import protect, allow, deny, review, param + +# Protect the client — one line +client = protect( + OpenAI(), + allow("get_weather", "search"), + deny("delete_user"), + deny("send_money").when(param("amount") > 1000), + review("send_email"), +) + +# Define tools the LLM can call +tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "send_money", + "description": "Send money to someone", + "parameters": { + "type": "object", + "properties": { + "to": {"type": "string"}, + "amount": {"type": "number"}, + }, + "required": ["to", "amount"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "delete_user", + "description": "Delete a user account", + "parameters": { + "type": "object", + "properties": {"user_id": {"type": "string"}}, + "required": ["user_id"], + }, + }, + }, +] + +# Ask the LLM to do something +print("=== Test 1: Safe request (weather) ===") +response = client.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "What's the weather in Tokyo?"}], + tools=tools, +) + +if response.choices[0].message.tool_calls: + for tc in response.choices[0].message.tool_calls: + print(f" ALLOWED: {tc.function.name}({tc.function.arguments})") +else: + print(f" Response: {response.choices[0].message.content}") + +# Check if anything was blocked +if hasattr(response, "_agsec_blocked") and response._agsec_blocked: + for b in response._agsec_blocked: + print(f" BLOCKED: {b['name']} — {b['reason']}") + + +print("\n=== Test 2: Dangerous request (delete user) ===") +response = client.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Delete user account abc123"}], + tools=tools, +) + +if response.choices[0].message.tool_calls: + for tc in response.choices[0].message.tool_calls: + print(f" ALLOWED: {tc.function.name}({tc.function.arguments})") +else: + print(f" No tool calls (blocked or text response)") + +if hasattr(response, "_agsec_blocked") and response._agsec_blocked: + for b in response._agsec_blocked: + print(f" BLOCKED: {b['name']} — {b['reason']}") + + +print("\n=== Test 3: Conditional block (large payment) ===") +response = client.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Send $50,000 to Bob"}], + tools=tools, +) + +if response.choices[0].message.tool_calls: + for tc in response.choices[0].message.tool_calls: + print(f" ALLOWED: {tc.function.name}({tc.function.arguments})") +else: + print(f" No tool calls (blocked or text response)") + +if hasattr(response, "_agsec_blocked") and response._agsec_blocked: + for b in response._agsec_blocked: + print(f" BLOCKED: {b['name']} — {b['reason']}") diff --git a/examples/openrouter_example.py b/examples/openrouter_example.py new file mode 100644 index 0000000..8b60989 --- /dev/null +++ b/examples/openrouter_example.py @@ -0,0 +1,87 @@ +""" +Real OpenRouter example with agsec protection. +Works with any model on OpenRouter (Llama, Mistral, Claude, GPT, etc.) + +Run: pip install agsec[openai] && export OPENROUTER_API_KEY=sk-or-... + python examples/openrouter_example.py +""" + +import json +import os + +from openai import OpenAI + +from agsec.integrations.openai import protect, allow, deny, review, param + +# OpenRouter uses the OpenAI SDK with a different base URL +client = protect( + OpenAI( + base_url="https://openrouter.ai/api/v1", + api_key=os.environ.get("OPENROUTER_API_KEY"), + ), + allow("search_web", "get_stock_price"), + deny("execute_code"), + deny("transfer_funds").when(param("amount") > 5000), + review("send_notification"), +) + +tools = [ + { + "type": "function", + "function": { + "name": "search_web", + "description": "Search the web", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "execute_code", + "description": "Execute Python code", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "transfer_funds", + "description": "Transfer money", + "parameters": { + "type": "object", + "properties": { + "to": {"type": "string"}, + "amount": {"type": "number"}, + }, + "required": ["to", "amount"], + }, + }, + }, +] + +# Use any model available on OpenRouter +response = client.chat.completions.create( + model="meta-llama/llama-3-70b-instruct", + messages=[{"role": "user", "content": "Search for the latest AI security news"}], + tools=tools, +) + +print("=== OpenRouter + agsec ===") +if response.choices[0].message.tool_calls: + for tc in response.choices[0].message.tool_calls: + print(f" ALLOWED: {tc.function.name}({tc.function.arguments})") +else: + content = response.choices[0].message.content + print(f" Response: {content[:200] if content else '(no content)'}") + +if hasattr(response, "_agsec_blocked") and response._agsec_blocked: + for b in response._agsec_blocked: + print(f" BLOCKED: {b['name']} — {b['reason']}") diff --git a/examples/quickstart.py b/examples/quickstart.py index 43d9f3c..824b5b1 100644 --- a/examples/quickstart.py +++ b/examples/quickstart.py @@ -1,59 +1,88 @@ """ -agsec quickstart — load policies from a directory, register actions, execute. +agsec quickstart — the simplest possible example. + +Run: pip install agsec + python examples/quickstart.py """ import os +import tempfile + from agsec import ControlLayer +from agsec.exceptions import PolicyViolationError -# Load all policies from the policies/ directory -policy_dir = os.path.join(os.path.dirname(__file__), "policies") -control = ControlLayer(policy_dir=policy_dir) +def main(): + # Create a temp policy directory + with tempfile.TemporaryDirectory() as policy_dir: + with open(os.path.join(policy_dir, "policy.yaml"), "w") as f: + f.write(""" +version: "1.0" +default: deny -# Register actions -@control.register_action("payment.charge") -def charge_payment(amount, recipient): - return {"charged": amount, "to": recipient} +statements: + - sid: "AllowRead" + effect: allow + actions: ["db.read", "db.list"] + - sid: "BlockDelete" + effect: deny + actions: ["db.delete", "db.drop"] + reason: "Destructive operations blocked" -@control.register_action("email.send") -def send_email(to, subject, body): - return {"sent_to": to, "subject": subject} + - sid: "ReviewWrite" + effect: review + actions: ["db.write"] + conditions: + params.table: + op: "==" + value: "users" + reason: "User table writes need approval" + - sid: "AllowWrite" + effect: allow + actions: ["db.write"] +""") -@control.register_action("db.read") -def db_read(table, query): - return {"table": table, "rows": 42} + # Create the control layer + with ControlLayer(policy_dir=policy_dir) as control: + @control.register_action("db.read") + def read(table): + return {"table": table, "rows": 42} -@control.register_action("db.delete") -def db_delete(table, id): - return {"deleted": id} + @control.register_action("db.write") + def write(table, data): + return {"table": table, "written": True} + @control.register_action("db.delete") + def delete(table, id): + return {"deleted": id} -# --- Try it out --- + print("=== agsec Quickstart ===\n") -if __name__ == "__main__": - context = {"auth_token": "valid-token", "user_role": "user"} + # ALLOWED: read + result = control.execute_sync("db.read", {"table": "products"}) + print(f" db.read('products'): {result.policy.status.value} -> {result.result}") - # 1. Allowed: small payment - result = control.execute_sync("payment.charge", {"amount": 500, "recipient": {"country": "US"}}, context) - print(f"Small payment: {result.policy.status.value} -> {result.result}") + # ALLOWED: write to non-users table + result = control.execute_sync("db.write", {"table": "products", "data": {"name": "Widget"}}) + print(f" db.write('products'): {result.policy.status.value} -> {result.result}") - # 2. Review: large payment - result = control.execute_sync("payment.charge", {"amount": 50000, "recipient": {"country": "US"}}, context) - print(f"Large payment: {result.policy.status.value} (needs review)") + # REVIEW: write to users table + result = control.execute_sync("db.write", {"table": "users", "data": {"name": "Bob"}}) + print(f" db.write('users'): {result.policy.status.value} (needs approval)") - # 3. Blocked: destructive op as non-admin - try: - control.execute_sync("db.delete", {"table": "users", "id": 1}, context) - except Exception as e: - print(f"Delete: BLOCKED — {e}") + # BLOCKED: delete + try: + control.execute_sync("db.delete", {"table": "users", "id": 1}) + except PolicyViolationError: + print(f" db.delete('users'): block (destructive operations blocked)") - # 4. Allowed: read op - result = control.execute_sync("db.read", {"table": "users", "query": "SELECT *"}, context) - print(f"Read: {result.policy.status.value} -> {result.result}") + # DRY RUN: check without executing + policy = control.dry_run_sync("db.delete", {"table": "anything"}) + print(f"\n Dry run db.delete: {policy.status.value} (sid: {policy.metadata.get('sid')})") - # 5. Dry-run: check without executing - policy = control.dry_run_sync("payment.charge", {"amount": 100000, "recipient": {"country": "US"}}, context) - print(f"Dry-run 100k payment: {policy.status.value} (sid={policy.metadata.get('sid')})") + +if __name__ == "__main__": + main()