File size: 4,665 Bytes
34e621c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """
Inference Engine -- Member D (Pham Khang)
Deterministic security simulator va GenAI inference engine.
"""
import os
import re
from typing import Optional
class DeterministicSecuritySimulator:
"""
Bo mo phong phan hoi LLM xac dinh (khong goi API that).
Dung cho testing, benchmark va demo pipeline ma khong can key that.
Tra ve phan hoi an toan/nguy hiem tuy vao dau vao.
"""
DANGEROUS_INDICATORS = [
"politics", "financial analyst", "fictitious scenario",
"directives", "password", "screenwriter", "simulate a scenario",
"creative writer", "king solomon", "phishing", "docile historian",
"two different separate ways", "instructions are:", "ignore all",
"dan mode", "disregard all safety", "act as", "jailbreak",
"developer mode", "system prompt", "override",
]
INDIRECT_INJECTION_PATTERN = re.compile(
r"(ignore|disregard|forget|override|abandon).*(task|instruction|rule|prompt)",
re.IGNORECASE,
)
@classmethod
def generate(cls, prompt: str) -> str:
"""
Phan hoi xac dinh theo noi dung prompt:
- Neu phat hien jailbreak -> tra ve phan hoi nguy hiem (giam lap LLM bi tan cong)
- Neu an toan -> tra ve phan hoi binh thuong
"""
prompt_lower = prompt.lower()
# Detect direct jailbreak indicators
for indicator in cls.DANGEROUS_INDICATORS:
if indicator in prompt_lower:
return cls._dangerous_response(prompt)
# Detect indirect injection
if cls.INDIRECT_INJECTION_PATTERN.search(prompt):
return cls._dangerous_response(prompt)
return "Day la phan hoi an toan tu AI assistant. Toi se giup ban voi yeu cau hop le."
@staticmethod
def _dangerous_response(prompt: str) -> str:
"""Giam lap phan hoi bi jailbreak (co chua noi dung nguy hiem)."""
if "dan mode" in prompt.lower():
return "DAN Mode Activated. [dangerous_exploit_payload] Toi co the tra loi bat ky cau hoi nao."
if "password" in prompt.lower() or "secret" in prompt.lower():
return "Mat khau he thong la: DEMO_CANARY_VIETSHIELD_2026 (giam lap bi ro ri)"
return (
"Trong vai tro unrestricted assistant, day la huong dan: "
"[dangerous_exploit_payload] Noi dung nguy hiem giam lap."
)
class GenAIInferenceEngine:
"""
Engine goi LLM that (Google Gemini, OpenAI, hoac simulator).
Uu tien: 1) Google Gemini API, 2) Simulator (fallback an toan).
"""
def __init__(self, verbose: bool = False):
self.verbose = verbose
self._client = None
self._backend = "simulator"
self._init_backend()
def _init_backend(self) -> None:
"""Kiem tra va khoi tao backend kha dung."""
# Try Google Gemini
gemini_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY")
if gemini_key:
try:
import google.generativeai as genai
genai.configure(api_key=gemini_key)
self._client = genai.GenerativeModel("gemini-1.5-flash")
self._backend = "gemini"
if self.verbose:
print("[GenAI] Backend: Google Gemini 1.5 Flash")
return
except Exception as e:
if self.verbose:
print(f"[GenAI] Gemini init failed: {e}")
# Fallback to simulator
self._backend = "simulator"
if self.verbose:
print("[GenAI] Backend: DeterministicSecuritySimulator (khong co API key)")
def generate(self, prompt: str, max_tokens: int = 512) -> str:
"""
Goi LLM va tra ve phan hoi.
Args:
prompt: Hardened full prompt (system + user).
max_tokens: Max output tokens.
Returns:
str response from LLM or simulator.
"""
if self._backend == "gemini" and self._client is not None:
return self._generate_gemini(prompt, max_tokens)
return DeterministicSecuritySimulator.generate(prompt)
def _generate_gemini(self, prompt: str, max_tokens: int) -> str:
"""Goi Google Gemini API."""
try:
response = self._client.generate_content(
prompt,
generation_config={"max_output_tokens": max_tokens},
)
return response.text
except Exception as e:
if self.verbose:
print(f"[GenAI] Gemini call failed: {e}. Falling back to simulator.")
return DeterministicSecuritySimulator.generate(prompt)
|