Custom Guardrail/Plugins Configuration

​

Template Repository Overview

​

Architecture

• main.py: FastAPI application with route definitions
• guardrail/: Directory containing all guardrail implementations
• entities.py: Pydantic models for request/response validation
• requirements.txt: Dependencies and libraries

​

Entities and Data Models

​

RequestContext

class SubjectType(str, Enum):
    user = 'user'
    team = 'team'
    serviceaccount = 'serviceaccount'

class Subject(BaseModel):
    subjectId: str
    subjectType: SubjectType
    subjectSlug: Optional[str] = None
    subjectDisplayName: Optional[str] = None

class RequestContext(BaseModel):
    user: Subject
    metadata: Optional[dict[str, str]] = None

​

InputGuardrailRequest

class InputGuardrailRequest(BaseModel):
    requestBody: CompletionCreateParams
    context: RequestContext
    config: Optional[dict] = None

​

OutputGuardrailRequest

class OutputGuardrailRequest(BaseModel):
    requestBody: CompletionCreateParams
    responseBody: ChatCompletion
    config: Optional[dict] = None
    context: RequestContext

​

Available Guardrails

from presidio_analyzer import AnalyzerEngine
from presidio_anonymizer import AnonymizerEngine

# Initialize Presidio
analyzer = AnalyzerEngine()
anonymizer = AnonymizerEngine()

def process_input_guardrail(request: InputGuardrailRequest) -> Optional[dict]:
    messages = request.requestBody.get('messages', [])
    for message in messages:
        # Analyze and anonymize each message
        results = analyzer.analyze(text=message["content"], entities=[], language='en')
        anonymized_content = anonymizer.anonymize(text=message["content"], analyzer_results=results)
        # ... update message content as needed
    # Return modified request if any PII was redacted, else None

from transformers import pipeline

classifier = pipeline("text-classification", model="unitary/unbiased-toxic-roberta")

def nsfw_filtering(request: OutputGuardrailRequest) -> Optional[dict]:
    transformed_body = request.responseBody.copy()
    for choice in transformed_body.get("choices", []):
        classification_results = classifier(choice["message"]["content"])
        for result in classification_results:
            if (
                (result['label'] == 'toxicity' and result['score'] > 0.2) or
                (result['label'] == 'sexual_explicit' and result['score'] > 0.2) or
                (result['label'] == 'obscene' and result['score'] > 0.2)
            ):
                raise HTTPException(
                    status_code=400, 
                    detail=f"This message is not allowed as it is NSFW"
                )
    return None

from guardrails import Guard
from guardrails.hub import MentionsDrugs

guard = Guard().use(MentionsDrugs, on_fail="exception")

def drug_mention(request: OutputGuardrailRequest) -> Optional[dict]:
    try:
        for choice in request.responseBody.get("choices", []):
            guard.validate(choice["message"]["content"])
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))
    return None

from guardrails import Guard
from guardrails_grhub_web_sanitization import WebSanitization

guard = Guard().use(WebSanitization, on_fail="exception")

def web_sanitization(request: InputGuardrailRequest) -> Optional[dict]:
    try:
        messages = request.requestBody.get("messages", [])
        for message in messages:
            guard.validate(message["content"])
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))
    return None

from guardrails import Guard
from guardrails.hub import DetectPII

guard = Guard().use(DetectPII, on_fail="exception")

def pii_detection_guardrails_ai(request: InputGuardrailRequest) -> Optional[dict]:
    try:
        messages = request.requestBody.get("messages", [])
        for message in messages:
            guard.validate(message["content"])
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))
    return None

​

Request Examples

{
  "requestBody": {
    "messages": [
      {
        "role": "user",
        "content": "Hello, my name is John Doe and my email is john.doe@example.com"
      }
    ],
    "model": "gpt-3.5-turbo",
    "temperature": 0.7
  },
  "config": {
    "check_content": true,
    "transform_input": true
  },
  "context": {
    "user": {
      "subjectId": "123",
      "subjectType": "user",
      "subjectSlug": "john_doe@truefoundry.com",
      "subjectDisplayName": "John Doe"
    },
    "metadata": {
      "ip_address": "192.168.1.1",
      "session_id": "abc123"
    }
  }
}

{
  "requestBody": {
    "messages": [
      {
        "role": "user",
        "content": "Hello"
      }
    ],
    "model": "gpt-3.5-turbo"
  },
  "responseBody": {
    "id": "chatcmpl-123",
    "object": "chat.completion",
    "created": 1677652288,
    "model": "gpt-3.5-turbo",
    "choices": [
      {
        "index": 0,
        "message": {
          "role": "assistant",
          "content": "Hello! How can I help you today?"
        },
        "finish_reason": "stop"
      }
    ]
  },
  "config": {
    "check_content": true
  },
  "context": {
    "user": {
      "subjectId": "123",
      "subjectType": "user",
      "subjectSlug": "john_doe@truefoundry.com",
      "subjectDisplayName": "John Doe"
    },
    "metadata": {
      "ip_address": "192.168.1.1",
      "session_id": "abc123"
    }
  }
}

​

Running Locally

# Install dependencies
pip install -r requirements.txt

# Run the server
python main.py

# Or using uvicorn directly
uvicorn main:app --host 0.0.0.0 --port 8000 --reload

​

Adding Custom Guardrail Integration

1. Navigate to AI Gateway
   • Go to AI Gateway in your TrueFoundry dashboard.
2. Access Guardrails
   • Click on Guardrails.
3. Add New Guardrails Group
   • Click on Add New Guardrails Group.

Navigate to Guardrails

4. Fill in the Guardrails Group Form
   • Name: Enter a name for your guardrails group.
   • Collaborators: Add collaborators who will have access to this group.
   • Custom Guardrail Config:
     • Name: Enter a name for the Custom Guardrail configuration.
     • Operation: The operation type to use for the guardrail.
       • Validate: Guardrails with this operation are used to validate requests. These guardrails are run in parallel.
       • Mutate: Guardrails with this operation can both validate and mutate requests. Mutate guardrails are run sequentially.
     • URL: Enter the URL for the Guardrail Server.
     • Auth Data: Provide authentication data for the Guardrail Server. This data will be sent to the Guardrail Server for authorization.
       • Choose between Custom Basic Auth or Custom Bearer Auth.
     • Headers (Optional): Add any headers required for the Guardrail Server. These will be forwarded as is.
     • Config: Enter the configuration for the Guardrail Server. This is a JSON object that will be sent along with the request.

Fill in the Custom Guardrail Form

​

How Custom Guardrail Config Relates to Guardrail Requests

• Config Propagation:
  The Config field you specify in the integration creation form is sent as the config attribute in every guardrail request payload. This allows you to parameterize your guardrail logic (e.g., set thresholds, enable/disable features, or pass secrets) without changing your server code.
• Request Structure:
  When a request is routed through a guardrail, the AI Gateway constructs a request object (such as InputGuardrailRequest or OutputGuardrailRequest) and sends it to your server. This object includes:
  • The original model input (requestBody)
  • (For output guardrails) The model’s response (responseBody)
  • The config object (from your integration creation form)
  • The context (user, metadata, etc.)
• Example Payload:
  Copy

  Ask AI

  {
    "requestBody": { /* original model input */ },
    "responseBody": { /* model output, for output guardrails */ },
    "config": { /* your custom config from the integration creation form */ },
    "context": { /* user and request metadata */ }
  }
  

• Dynamic Behavior:
  By updating the Custom Guardrail Config in the integration creation form, you can change the behavior of your guardrail server in real time—no code redeploy required. For example, you might adjust PII detection sensitivity, toggle logging, or update allowed user lists.

​

Example: Sending a Request to Your Guardrail Server

​

Sample Input Guardrail Request Payload & cURL Example

​

Sample Output Guardrail Request Payload & cURL Example