feat(mock): llm-katan support

Signed-off-by: FeiDaLI <[email protected]>

Author: FeiDaLI

.pre-commit-config.yaml

@@ -3,7 +3,7 @@
 repos:
 # Basic hooks for Go, Rust, Python And JavaScript files only
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v6.0.0
+  rev: v4.4.0
   hooks:
   - id: trailing-whitespace
     files: \.(go|rs|py|js)$
@@ -71,15 +71,15 @@ repos:
 
 # Python specific hooks
 - repo: https://github.com/psf/black
-  rev: 25.1.0
+  rev: 22.12.0
   hooks:
   - id: black
     language_version: python3
     files: \.py$
     exclude: ^(\.venv/|venv/|env/|__pycache__/|\.git/|site/)
 
 - repo: https://github.com/PyCQA/isort
-  rev: 6.0.1
+  rev: 5.12.0
   hooks:
   - id: isort
     args: ["--profile", "black"]

e2e-tests/06-pii-detection-test.py

@@ -652,7 +652,9 @@ def test_extproc_pii_detection_comprehensive(self):
             status = (
                 "🔒"
                 if result["is_blocked"]
-                else "✅" if result["request_allowed"] else "❌"
+                else "✅"
+                if result["request_allowed"]
+                else "❌"
             )
             print(f"  {status} {result['test_case']}")
             print(f"      Content: {result['content']}")
@@ -671,9 +673,7 @@ def test_extproc_pii_detection_comprehensive(self):
                 passed=False,
                 message="⚠️ No clear evidence of ExtProc PII detection in production pipeline",
             )
-            print(
-                "📝 NOTE: This may indicate PII detection is not active in ExtProc or"
-            )
+            print("📝 NOTE: This may indicate PII detection is not active in ExtProc or")
             print("         PII policies are configured to allow all content through")
 
     def test_multiple_pii_types_analysis(self):

src/training/classifier_model_fine_tuning/ft_linear.py

@@ -508,19 +508,25 @@ def tokenize_function(examples):
 
     # Evaluate on validation set
     logger.info("Evaluating on validation set...")
-    val_accuracy, val_report, val_conf_matrix, val_predictions = (
-        evaluate_category_classifier(
-            model, tokenizer, val_texts, val_categories, idx_to_category, device
-        )
+    (
+        val_accuracy,
+        val_report,
+        val_conf_matrix,
+        val_predictions,
+    ) = evaluate_category_classifier(
+        model, tokenizer, val_texts, val_categories, idx_to_category, device
     )
     logger.info(f"Validation accuracy: {val_accuracy:.4f}")
 
     # Evaluate on test set
     logger.info("Evaluating on test set...")
-    test_accuracy, test_report, test_conf_matrix, test_predictions = (
-        evaluate_category_classifier(
-            model, tokenizer, test_texts, test_categories, idx_to_category, device
-        )
+    (
+        test_accuracy,
+        test_report,
+        test_conf_matrix,
+        test_predictions,
+    ) = evaluate_category_classifier(
+        model, tokenizer, test_texts, test_categories, idx_to_category, device
     )
     logger.info(f"Test accuracy: {test_accuracy:.4f}")
 

src/training/prompt_guard_fine_tuning/jailbreak_bert_finetuning.py

@@ -100,9 +100,9 @@
 import torch
 
 # Suppress common non-critical warnings
-os.environ["TOKENIZERS_PARALLELISM"] = (
-    "false"  # Suppress tokenizer parallelism warnings
-)
+os.environ[
+    "TOKENIZERS_PARALLELISM"
+] = "false"  # Suppress tokenizer parallelism warnings
 warnings.filterwarnings(
     "ignore", message=".*TensorFloat32.*"
 )  # Suppress TF32 performance hints
@@ -2343,19 +2343,25 @@ def tokenize_function(examples):
 
     # Evaluate on validation set
     logger.info("Evaluating on validation set...")
-    val_accuracy, val_report, val_conf_matrix, val_predictions = (
-        evaluate_jailbreak_classifier(
-            model, tokenizer, val_texts, val_categories, idx_to_category, device
-        )
+    (
+        val_accuracy,
+        val_report,
+        val_conf_matrix,
+        val_predictions,
+    ) = evaluate_jailbreak_classifier(
+        model, tokenizer, val_texts, val_categories, idx_to_category, device
     )
     logger.info(f"Validation accuracy: {val_accuracy:.4f}")
 
     # Evaluate on test set
     logger.info("Evaluating on test set...")
-    test_accuracy, test_report, test_conf_matrix, test_predictions = (
-        evaluate_jailbreak_classifier(
-            model, tokenizer, test_texts, test_categories, idx_to_category, device
-        )
+    (
+        test_accuracy,
+        test_report,
+        test_conf_matrix,
+        test_predictions,
+    ) = evaluate_jailbreak_classifier(
+        model, tokenizer, test_texts, test_categories, idx_to_category, device
     )
     logger.info(f"Test accuracy: {test_accuracy:.4f}")
 

src/training/training_lora/classifier_model_fine_tuning_lora/ft_linear_lora.py

@@ -69,13 +69,7 @@
 import torch
 import torch.nn as nn
 from datasets import Dataset, load_dataset
-from peft import (
-    LoraConfig,
-    PeftConfig,
-    PeftModel,
-    TaskType,
-    get_peft_model,
-)
+from peft import LoraConfig, PeftConfig, PeftModel, TaskType, get_peft_model
 from sklearn.metrics import accuracy_score, f1_score, precision_recall_fscore_support
 from sklearn.model_selection import train_test_split
 from transformers import (

src/training/training_lora/pii_model_fine_tuning_lora/pii_bert_finetuning_lora.py

@@ -70,13 +70,7 @@
 import torch
 import torch.nn as nn
 from datasets import Dataset, load_dataset
-from peft import (
-    LoraConfig,
-    PeftConfig,
-    PeftModel,
-    TaskType,
-    get_peft_model,
-)
+from peft import LoraConfig, PeftConfig, PeftModel, TaskType, get_peft_model
 from sklearn.metrics import accuracy_score, f1_score, precision_recall_fscore_support
 from sklearn.model_selection import train_test_split
 from transformers import (

src/training/training_lora/prompt_guard_fine_tuning_lora/jailbreak_bert_finetuning_lora.py

@@ -77,13 +77,7 @@
 import torch
 import torch.nn as nn
 from datasets import Dataset, load_dataset
-from peft import (
-    LoraConfig,
-    PeftConfig,
-    PeftModel,
-    TaskType,
-    get_peft_model,
-)
+from peft import LoraConfig, PeftConfig, PeftModel, TaskType, get_peft_model
 from sklearn.metrics import accuracy_score, f1_score, precision_recall_fscore_support
 from sklearn.model_selection import train_test_split
 from transformers import (

tools/llm-katan-server/README.md

@@ -5,6 +5,7 @@ A FastAPI wrapper around [llm-katan](https://pypi.org/project/llm-katan/) that p
 ## Architecture
 
 This server acts as a proxy that:
+
 1. Receives OpenAI-compatible API requests
 2. Forwards them to a running `llm-katan` instance
 3. Returns the responses with proper model name mapping
@@ -78,24 +79,3 @@ services:
     depends_on:
       - llm-katan-backend
 ```
-
-### Testing
-
-```bash
-# Health check
-curl http://localhost:8000/health
-
-# List models
-curl http://localhost:8000/v1/models
-
-# Chat completion (uses real LLM)
-curl -X POST http://localhost:8000/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "model": "Qwen/Qwen2-0.5B-Instruct",
-    "messages": [{"role": "user", "content": "Hello!"}],
-    "max_tokens": 50
-  }'
-```
-
-Intended for local testing with Docker Compose profile `testing`.

tools/llm-katan-server/app.py

@@ -1,9 +1,9 @@
 import math
-import time
 import os
-import requests
+import time
 from typing import List, Optional
 
+import requests
 import uvicorn
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
@@ -18,7 +18,9 @@
 # Check if HuggingFace token is set
 hf_token = os.getenv("HUGGINGFACE_HUB_TOKEN")
 if not hf_token:
-    print("Warning: HUGGINGFACE_HUB_TOKEN not set. Some models may require authentication.")
+    print(
+        "Warning: HUGGINGFACE_HUB_TOKEN not set. Some models may require authentication."
+    )
 
 
 class ChatMessage(BaseModel):
@@ -49,37 +51,37 @@ async def chat_completions(req: ChatRequest):
         # Forward request to llm-katan backend
         llm_katan_request = {
             "model": MODEL,
-            "messages": [{"role": msg.role, "content": msg.content} for msg in req.messages],
+            "messages": [
+                {"role": msg.role, "content": msg.content} for msg in req.messages
+            ],
             "temperature": req.temperature,
         }
-        
+
         if req.max_tokens:
             llm_katan_request["max_tokens"] = req.max_tokens
-        
+
         # Make request to llm-katan
         response = requests.post(
-            f"{LLM_KATAN_URL}/v1/chat/completions",
-            json=llm_katan_request,
-            timeout=30
+            f"{LLM_KATAN_URL}/v1/chat/completions", json=llm_katan_request, timeout=30
         )
-        
+
         if response.status_code != 200:
             raise HTTPException(
                 status_code=response.status_code,
-                detail=f"LLM Katan error: {response.text}"
+                detail=f"LLM Katan error: {response.text}",
             )
-        
+
         result = response.json()
-        
+
         # Update the model name in response to match our served model name
         result["model"] = req.model
-        
+
         return result
-        
+
     except requests.exceptions.RequestException as e:
         # Fallback to simple echo behavior if llm-katan is not available
         print(f"Warning: LLM Katan not available ({e}), using fallback response")
-        
+
         # Simple echo-like behavior as fallback
         last_user = next(
             (m.content for m in reversed(req.messages) if m.role == "user"), ""

website/docs/installation/installation.md

@@ -126,12 +126,14 @@ model_config:
 If you don't have a vLLM backend set up, you can use the provided mock services for testing:
 
 **Option 1: Mock vLLM (Simple Echo Service)**
+
 ```bash
 # Start a simple mock service that echoes back responses
 python tools/mock-vllm/app.py
 ```
 
 **Option 2: LLM Katan Server (Real LLM with Lightweight Backend)**
+
 ```bash
 # First, start llm-katan backend (requires pip install llm-katan)
 llm-katan --model Qwen/Qwen2-0.5B-Instruct --port 8001
@@ -141,6 +143,7 @@ python tools/llm-katan-server/app.py
 ```
 
 For the mock services, update your `config/config.yaml`:
+
 ```yaml
 vllm_endpoints:
   - name: "mock-endpoint"
@@ -152,6 +155,7 @@ vllm_endpoints:
       - "Qwen/Qwen2-0.5B-Instruct"  # For llm-katan-server
     weight: 1
 ```
+
 :::
 
 :::note[**Important: Address Format Requirements**]