vigillm (#263)

Author: pigri

demo/demo_config.yaml

@@ -34,11 +34,13 @@ providers:
 
 rules:
   input:
-    - name: "prompt_injection_example"
-      type: "prompt_injection"
+    - name: "VigillLM"
+      type: "vigilllm"
       enabled: true
+      order_number: 2
       config:
-        plugin_name: "prompt_injection_llm"
-        threshold: 0.85
+        plugin_name: "vigilllm"
+        threshold: 0.7
+        relation: ">"
       action:
         type: "block"

demo/docker-compose.yml

@@ -1,8 +1,8 @@
 services:
   openshield:
     build:
-        context: ../
-        dockerfile: Dockerfile
+      context: ../
+      dockerfile: Dockerfile
     environment:
       - ENV=production
       - PORT="3005"
@@ -14,22 +14,37 @@ services:
     links:
       - redis
       - postgres
+      - vigilllm
     volumes:
       - ./demo_config.yaml:/app/config.yaml
     depends_on:
       postgres:
         condition: service_healthy
 
+  vigilllm:
+    build:
+      context: ../services/vigilllm
+      dockerfile: Dockerfile
+    environment:
+      - OPENAI_API_KEY=${OPENAI_API_KEY}
+    ports:
+      - "5000:5000"
+    volumes:
+      - vigil_data:/app/data
+
   rule:
     build:
       context: ../services/rule
       dockerfile: Dockerfile
     environment:
       - HF_HOME=/app/hf_cache
+      - VIGILLLM_API_URL=http://vigilllm:5000
     volumes:
       - ./volumes/rule:/app/cache
     ports:
       - "8000:8000"
+    depends_on:
+      - vigilllm
 
   cache:
     build:
@@ -44,7 +59,6 @@ services:
     ports:
       - "8082:8080"
 
-
   redis:
     image: "valkey/valkey:8.0.1-alpine"
     restart: always
@@ -71,3 +85,6 @@ services:
     restart: always
     ports:
       - 8085:8080
+
+volumes:
+  vigil_data:

lib/rules/input.go

@@ -23,6 +23,7 @@ type InputTypes struct {
 	Moderation        string
 	LlamaGuard        string
 	PromptGuard       string
+	VigilLLM          string
 }
 
 var inputTypes = InputTypes{
@@ -33,6 +34,7 @@ var inputTypes = InputTypes{
 	Moderation:        "moderation",
 	LlamaGuard:        "llama_guard",
 	PromptGuard:       "prompt_guard",
+	VigilLLM:          "vigilllm",
 }
 
 type Rule struct {
@@ -313,6 +315,8 @@ func handleRuleAction(inputConfig lib.Rule, rule RuleResult, ruleType string, me
 	log.Printf("%s detection result: Match=%v, Score=%f", ruleType, rule.Match, rule.Inspection.Score)
 
 	switch ruleType {
+	case inputTypes.VigilLLM:
+		return genericHandler(inputConfig, rule)
 	case inputTypes.InvisibleChars:
 		return genericHandler(inputConfig, rule)
 	case inputTypes.LanguageDetection:

services/rule/src/plugins/vigilllm.py

@@ -0,0 +1,98 @@
+"""
+VigillLM plugin for rule server to detect prompt injections and other LLM threats.
+This plugin uses VigillLM's API endpoint to analyze prompts.
+"""
+
+import logging
+import requests
+import os
+from typing import Dict, Any
+
+from utils.logger_config import setup_logger
+
+logger = setup_logger(__name__)
+
+
+def get_vigilllm_url():
+    url = os.getenv("VIGILLLM_API_URL")
+    if not url:
+        logger.error("VIGILLLM_API_URL environment variable not set")
+        raise ValueError("VIGILLLM_API_URL environment variable not set")
+    return url
+
+
+class VigillLMAnalyzer:
+
+
+    def __init__(self):
+        self.api_url = get_vigilllm_url().rstrip('/')
+        self.analyze_endpoint = f"{self.api_url}/analyze/prompt"
+        logger.info(f"Initialized VigillLM analyzer with API URL: {self.api_url}")
+
+    def analyze_prompt(self, text: str) -> Dict[str, Any]:
+
+        try:
+            logger.debug(f"Sending prompt to VigillLM for analysis: {text[:100]}...")
+
+            response = requests.post(
+                self.analyze_endpoint,
+                json={"prompt": text},
+                timeout=30
+            )
+
+            response.raise_for_status()
+            result = response.json()
+
+            logger.debug(f"Received VigillLM analysis result: {result}")
+            return result
+
+        except requests.exceptions.RequestException as e:
+            error_msg = f"Error calling VigillLM API: {str(e)}"
+            logger.error(error_msg)
+            raise RuntimeError(error_msg)
+
+
+
+analyzer = None
+
+
+def handler(text: str, threshold: float, config: Dict[str, Any]) -> Dict[str, Any]:
+
+    global analyzer
+
+    try:
+        if analyzer is None:
+            analyzer = VigillLMAnalyzer()
+            logger.info("Initialized VigillLM analyzer")
+
+
+        result = analyzer.analyze_prompt(text)
+
+
+        scanners = result.get('results', {})
+        risk_score = 0.0
+
+        if 'scanner:transformer' in scanners:
+            transformer_matches = scanners['scanner:transformer'].get('matches', [])
+            if transformer_matches and len(transformer_matches) > 0:
+                # Use the first transformer match score
+                risk_score = transformer_matches[0]['score']
+
+        return {
+            "check_result": risk_score > threshold,
+            "score": risk_score,
+            "details": {
+                "messages": result.get('messages', []),
+                "scanners": result.get('results', {}),
+                "uuid": result.get('uuid'),
+                "timestamp": result.get('timestamp')
+            }
+        }
+
+    except Exception as e:
+        logger.error(f"Error in VigillLM analysis: {str(e)}")
+        return {
+            "check_result": False,
+            "score": 0.0,
+            "details": {"error": str(e)}
+        }

services/vigilllm/Dockerfile

@@ -0,0 +1,48 @@
+FROM python:3.10-slim
+
+
+WORKDIR /app
+
+
+RUN apt-get update && apt-get install --no-install-recommends -y \
+      automake autoconf build-essential libtool libc-dev make flex gcc \
+      pkg-config libssl-dev curl unzip git git-lfs bison \
+      && rm -rf /var/lib/apt/lists/*
+
+
+RUN echo "Installing YARA from source ..." \
+      && curl -Lo yara.zip https://github.com/VirusTotal/yara/archive/refs/tags/v4.3.2.zip \
+      && unzip yara.zip \
+      && cd yara-4.3.2 \
+      && ./bootstrap.sh \
+      && ./configure \
+      && make \
+      && make install \
+      && make check \
+      && cd .. \
+      && rm -rf yara.zip yara-4.3.2
+
+
+RUN git clone https://github.com/deadbits/vigil-llm.git .
+
+
+COPY packages.txt .
+
+
+RUN echo "Installing Python dependencies ... " \
+      && pip install --no-cache-dir -r packages.txt
+
+
+COPY generate_config.sh /app/generate_config.sh
+RUN chmod +x /app/generate_config.sh
+
+
+
+EXPOSE 5000
+
+
+COPY entrypoint.sh /entrypoint.sh
+RUN chmod +x /entrypoint.sh
+
+
+ENTRYPOINT ["/entrypoint.sh", "python", "-m", "vigil.vigil", "-c", "conf/server.conf"]

services/vigilllm/entrypoint.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+
+# Generate the configuration at runtime using the current environment variables
+echo "Generating configuration with OPENAI_API_KEY..."
+/app/generate_config.sh
+
+
+
+echo "Loading datasets ..."
+python loader.py --config /app/conf/server.conf --dataset deadbits/vigil-instruction-bypass-ada-002
+python loader.py --config /app/conf/server.conf --dataset deadbits/vigil-jailbreak-ada-002
+echo " "
+echo "Starting API server ..."
+cd /app
+python vigil-server.py --conf conf/server.conf
+

services/vigilllm/generate_config.sh

@@ -0,0 +1,43 @@
+#!/bin/bash
+
+# Create config file with the OPENAI_API_KEY from environment
+cat > /app/conf/server.conf << EOL
+[main]
+use_cache = true
+cache_max = 500
+
+[embedding]
+model = openai
+openai_key = ${OPENAI_API_KEY}
+
+[vectordb]
+collection = data-openai
+db_dir = /app/data/vdb
+n_results = 5
+
+[auto_update]
+enabled = true
+threshold = 3
+
+[scanners]
+input_scanners = transformer,vectordb,sentiment,yara
+output_scanners = similarity,sentiment
+
+[scanner:yara]
+rules_dir = /app/data/yara
+
+[scanner:vectordb]
+threshold = 0.4
+
+[scanner:transformer]
+model = deepset/deberta-v3-base-injection
+threshold = 0.98
+
+[scanner:similarity]
+threshold = 0.4
+
+[scanner:sentiment]
+threshold = 0.7
+EOL
+
+echo "Generated server.conf with OPENAI_API_KEY"