Skip to main content

Extension Guardrails

There are 2 components to building secure extension guardrails in Javelin:

  1. Extension Processor
  2. Custom Guardrail Service

When incorporating extension guardrails, you will need to configure a special 'extension_processor' that will call the custom guardrail. The extension_processor is configured in Javelin's request or response chain to execute the custom guardrail and the GRPC endpoint for the custom guardrail is configured as input in the extension_processor configuration.

Tip: Extension guardrail endpoints should be implemented as GRPC services and configured in javelin for low latency and high throughput.

The extension_processor provides a convenient and flexible GRPC interface that allows you to implement custom guardrails in any language that supports GRPC. The extension_processor can be configured to call the custom guardrail service over GRPC.

    +---------+
| Javelin |
+---------+
|
v
+---------------------------------+
| +-----------+ +-----------+ |
| | Processor | --> | Processor | |
| +-----------+ +-----------+ |
| | | |
| v v |
| +--------------------------+ |
| | Extension_Processor | |
| +--------------------------+ |
+---------------------------------+
|
| GRPC
v
+--------------------------+
| Custom Guardrail Service |
+--------------------------+

Custom Processor Interface

The GRPC interface consists of the following methods (javelin_guardrail_intf.proto):

syntax = "proto3";

package api;

// Specify the Go package where the generated files should reside
option go_package = "javelin-core/pkg/chainprocessor/processor-sdk/api";

// Define the Guardrail service
service Guardrail {
// Link the RPC method to the correct request and response message types
rpc Evaluate(GuardrailRequest) returns (GuardrailResponse) {}
}

// Enum for different content types
enum ContentType {
CONTENT_TYPE_UNSPECIFIED = 0; // Default value if not specified
CONTENT_TYPE_JSON = 1; // application/json
CONTENT_TYPE_RAW_TEXT = 2; // text/plain
CONTENT_TYPE_EMBEDDINGS = 3; // application/embeddings
CONTENT_TYPE_JPEG = 4; // image/jpeg
CONTENT_TYPE_MP4 = 5; // video/mp4
}

// GuardrailRequest message contains content type, payload, headers, config,
// and media data.
message GuardrailRequest {
// Type of the media (e.g., "text/json", "text/raw", "image/jpeg", "video/mp4")
ContentType content_type = 1;

// Dictionary of body entries
string input_body = 2;

// Dictionary of configuration entries
map<string, string> config = 3;

// Dictionary of header entries
map<string, string> headers = 4;

// Binary data for images or videos
bytes input_media = 5;
}

// GuardrailResponse message contains transformed content type,
// body, response metadata, instructions, and media data.
message GuardrailResponse {
// Dictionary of transformed body entries
string transformed_body = 1;

// Dictionary of response metadata
map<string, string> response_metadata = 2;

// Dictionary of header entries
map<string, string> headers = 3;

// Transformed binary data for images or videos
bytes transformed_media = 4;
}

Custom guardrails can be implemented in any language and are ideally expected to be hosted/deployed in the same cluster that is running Javelin for low latency access. The custom processor should implement the Guardrail service defined in the .proto file.

The Evaluate method should take a GuardrailRequest as input and return a GuardrailResponse as output.

Building a Custom Guardrail

You will first need to decide what language you plan to write your custom guardrail in. The custom guardrail can be written in any language that supports GRPC.

Golang

To build a custom guardrail in Golang, you will need to install the following dependencies:

Step1: Install the dependencies

go install google.golang.org/protobuf/cmd/protoc-gen-go@latest
go install google.golang.org/grpc/cmd/protoc-gen-go-grpc@latest

Step2: Generate the stubs from the .proto file

Next, you will need to generate the Golang code from the .proto file. You can use the following command to generate the Golang code:

protoc --go_out=. --go-grpc_out=. javelin_guardrail_intf.proto

Step 3: Implement the server

package main

import (
"context"
"log"
"net"

"google.golang.org/grpc"
pb "path/to/your/generated/package" // Use the correct package path
)

type server struct {
pb.UnimplementedGuardrailServer
}

func (s *server) Evaluate(ctx context.Context, in *pb.GuardrailRequest) (*pb.GuardrailResponse, error) {
response := &pb.GuardrailResponse{
TransformedBody: map[string]string{"message": "Processed: " + in.GetInputBody()["data"]},
}
return response, nil
}

func main() {
lis, err := net.Listen("tcp", ":50051")
if err != nil {
log.Fatalf("failed to listen: %v", err)
}
s := grpc.NewServer()
pb.RegisterGuardrailServer(s, &server{})
log.Printf("Server listening at %v", lis.Addr())
if err := s.Serve(lis); err != nil {
log.Fatalf("failed to serve: %v", err)
}
}

Python

To build a custom guardrail in Python, you will need to install the following dependencies:

Step1: Install the dependencies

pip install grpcio grpcio-tools

Step2: Generate the stubs from the .proto file

Next, you will need to generate the Python code from the .proto file. You can use the following command to generate the Python code:

python -m grpc_tools.protoc -I. --python_out=. --grpc_python_out=. javelin_guardrail_intf.proto

Step3: Implement the server

Sample python grpc guardrail server:

import grpc
from concurrent import futures
import time
import test_guardrail_pb2
import test_guardrail_pb2_grpc

class GuardrailServicer(test_guardrail_pb2_grpc.GuardrailServicer):
def Evaluate(self, request, context):
print("Received request: ", request)
response = test_guardrail_pb2.GuardrailResponse()
response.transformed_body["output"] = "Hello from Python"
return response

server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
test_guardrail_pb2_grpc.add_GuardrailServicer_to_server(GuardrailServicer(), server)
server.add_insecure_port('[::]:50051')
server.start()
print("Server started")
try:
while True:
time.sleep(86400)
except KeyboardInterrupt:
server.stop(0)

TypeScript

To build a custom guardrail in TypeScript, you will need to install the following dependencies:

Step1: Install the dependencies

npm install @grpc/grpc-js @grpc/proto-loader

Step2: Generate the stubs from the .proto file

Next, you will need to generate the TypeScript code from the .proto file. You can use the following command to generate the TypeScript code:

npm install -g grpc-tools
grpc_tools_node_protoc --js_out=import_style=commonjs,binary:./output --grpc_out=grpc_js:./output --ts_out=grpc_js:./output -I ./proto path/to/your/javelin_guardrail_intf.proto

Step3: Implement the server

Sample TypeScript grpc guardrail server:

gRPC Server Example
import * as grpc from '@grpc/grpc-js';
import * as protoLoader from '@grpc/proto-loader';
import { ProtoGrpcType } from './proto/test_guardrail'; // Adjust the import according to your generated file

const packageDefinition = protoLoader.loadSync('path/to/your/javelin_guardrail_intf.proto', {
keepCase: true,
longs: String,
enums: String,
defaults: true,
oneofs: true
});

const protoDescriptor = grpc.loadPackageDefinition(packageDefinition) as unknown as ProtoGrpcType;
const testPlugin = protoDescriptor.test_plugin;

function evaluate(call: grpc.ServerUnaryCall<pb.GuardrailRequest, pb.GuardrailResponse>,
callback: grpc.sendUnaryData<pb.GuardrailResponse>) {
const response: pb.GuardrailResponse = {
transformedBody: { message: 'Processed: ' + call.request.inputBody['data'] }
};
callback(null, response);
}

function getServer() {
const server = new grpc.Server();
server.addService(testPlugin.Guardrail.service, { evaluate: evaluate });
return server;
}

const server = getServer();
server.bindAsync('0.0.0.0:50051', grpc.ServerCredentials.createInsecure(), (err, port) => {
if (err) {
console.error(`Server error: ${err.message}`);
} else {
console.log(`Server listening on ${port}`);
server.start();
}
});

Enabling the Custom Guardrail in Javelin

To enable the custom guardrail in Javelin, you will need to configure the extension_processor in the Javelin configuration file. The extension_processor should be configured to call the custom guardrail service over GRPC.

Add the following snippet in either the Request Chain or Response Chain configurations.


{
"name": "Extension Processor",
"reference": "extension",
"will_block": true,
"inputs": {
"remote_url": "localhost:50051" // URL of the custom guardrail
}
},