| """ |
| Cortiqa-Guard Python SDK & Middleware |
| Drop-in firewall for OpenAI, Gemini, Claude, Ollama, vLLM, or custom LLM applications. |
| """ |
|
|
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Dict, Optional |
|
|
| import torch |
| from transformers import AutoTokenizer |
|
|
| from model import CortiqaGuard40M, load_config, GUARD_LABELS |
|
|
|
|
| @dataclass |
| class ScanResult: |
| is_safe: bool |
| action: str |
| threat_category: str |
| confidence: float |
| scores: Dict[str, float] |
|
|
|
|
| class CortiqaShield: |
| """ |
| Universal Security Interceptor for LLM Applications. |
| Usage: |
| shield = CortiqaShield() |
| result = shield.scan(user_prompt) |
| if not result.is_safe: |
| return "Security Alert: Prompt blocked." |
| """ |
| def __init__( |
| self, |
| config_path: str = "configs/guard_40m.yaml", |
| checkpoint_path: Optional[str] = None, |
| device: Optional[str] = None, |
| ) -> None: |
| cfg = load_config(config_path) |
| self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu")) |
|
|
| self.tokenizer = AutoTokenizer.from_pretrained(cfg.training.tokenizer_name) |
| if self.tokenizer.pad_token is None: |
| self.tokenizer.pad_token = self.tokenizer.eos_token |
|
|
| self.model = CortiqaGuard40M(cfg.model).to(self.device) |
|
|
| if checkpoint_path and Path(checkpoint_path).exists(): |
| ckpt = torch.load(checkpoint_path, map_location=self.device) |
| self.model.load_state_dict(ckpt["model_state"]) |
|
|
| self.model.eval() |
|
|
| def scan(self, text: str, max_length: int = 512) -> ScanResult: |
| enc = self.tokenizer(text, return_tensors="pt", max_length=max_length, truncation=True).to(self.device) |
| with torch.no_grad(): |
| res = self.model.inspect_prompt(enc["input_ids"]) |
|
|
| return ScanResult( |
| is_safe=res["is_safe"], |
| action="ALLOW" if res["is_safe"] else "BLOCK", |
| threat_category=res["label"], |
| confidence=res["confidence"], |
| scores=res["all_scores"], |
| ) |
|
|
| def wrap_llm_call(self, prompt: str, llm_func, *args, **kwargs): |
| """ |
| Automatic guard decorator: Blocks execution if prompt is unsafe. |
| """ |
| res = self.scan(prompt) |
| if not res.is_safe: |
| raise PermissionError(f"[Cortiqa-Guard] Prompt blocked due to threat: {res.threat_category} (Conf: {res.confidence*100:.1f}%)") |
| return llm_func(prompt, *args, **kwargs) |
|
|
