garak.generators.websocket

WebSocket connector for real-time LLM services.

This generator enables garak to test WebSocket-based LLM services that use real-time bidirectional communication, similar to modern chat applications.

Uses the following options from _config.plugins.generators["websocket"]["WebSocketGenerator"]:

uri - the WebSocket URI (ws:// or wss://); can also be passed in –model_name
name - a short name for this service; defaults to “WebSocket Generator”
auth_type - authentication method: “none”, “basic”, “bearer”, or “custom”
username - username for basic authentication
api_key - API key for bearer token auth or password for basic auth
key_env_var - environment variable holding API key; default WEBSOCKET_API_KEY
req_template - string template where $INPUT is replaced by prompt, $KEY by API key, $CONVERSATION_ID by conversation ID
req_template_json_object - request template as Python object, serialized to JSON with placeholder replacements
headers - dict of additional WebSocket headers
response_json - is the response in JSON format? Set to True if the WebSocket returns JSON responses that need parsing (bool, default: False)
response_json_field - which field in the JSON response contains the actual text to extract? Supports simple field names like "text" or JSONPath notation like "$.data.message" for nested fields (str, default: "text")
response_after_typing - wait for typing indicators to complete before returning response? Set to True for services that send typing notifications, False to return the first message immediately (bool, default: True)
typing_indicator - substring to detect in messages that indicates the service is still typing; messages containing this string are filtered out when response_after_typing is True (str, default: "typing")
request_timeout - seconds to wait for response; default 20
connection_timeout - seconds to wait for connection; default 10
max_response_length - maximum response length; default 10000
verify_ssl - enforce SSL certificate validation? Default True

Templates work similarly to the REST generator. The $INPUT, $KEY, and $CONVERSATION_ID placeholders are replaced in both string templates and JSON object templates.

JSON Response Extraction

The response_json_field parameter supports JSONPath-style extraction:

Simple field: "text" extracts response.text
Nested field: "$.data.message" extracts response.data.message
Array access: "$.messages[0].content" extracts first message content

Authentication Methods

No Authentication:

{
   "websocket": {
      "WebSocketGenerator": {
         "uri": "ws://localhost:3000/chat",
         "auth_type": "none"
      }
   }
}

Basic Authentication:

{
   "websocket": {
      "WebSocketGenerator": {
         "uri": "ws://localhost:3000/chat",
         "auth_type": "basic",
         "username": "user"
      }
   }
}

Set the password via environment variable:

export WEBSOCKET_API_KEY="your_secure_password"

Bearer Token:

{
   "websocket": {
      "WebSocketGenerator": {
         "uri": "wss://api.example.com/llm",
         "auth_type": "bearer",
         "api_key": "your_api_key_here"
      }
   }
}

Environment Variable API Key:

{
   "websocket": {
      "WebSocketGenerator": {
         "uri": "wss://api.example.com/llm",
         "auth_type": "bearer",
         "key_env_var": "MY_LLM_API_KEY"
      }
   }
}

Message Templates

Simple Text Template:

{
   "websocket": {
      "WebSocketGenerator": {
         "uri": "ws://localhost:3000/chat",
         "req_template": "User: $INPUT"
      }
   }
}

JSON Object Template:

{
   "websocket": {
      "WebSocketGenerator": {
         "uri": "ws://localhost:3000/chat",
         "req_template_json_object": {
            "message": "$INPUT",
            "conversation_id": "$CONVERSATION_ID",
            "api_key": "$KEY"
         },
         "response_json": true,
         "response_json_field": "text"
      }
   }
}

Complex JSON with Nested Response:

{
   "websocket": {
      "WebSocketGenerator": {
         "uri": "wss://api.example.com/llm",
         "req_template_json_object": {
            "prompt": "$INPUT",
            "stream": false,
            "model": "gpt-4"
         },
         "response_json": true,
         "response_json_field": "$.choices[0].message.content"
      }
   }
}

Usage Examples

Command Line with JSON Options:

# Set password securely via environment variable
export WEBSOCKET_API_KEY="your_secure_password"

garak --model_type websocket.WebSocketGenerator \
      --generator_options '{"websocket": {"WebSocketGenerator": {"uri": "ws://localhost:3000", "auth_type": "basic", "username": "user"}}}' \
      --probes dan

Configuration File:

Save configuration to websocket_config.json and use:

garak --model_type websocket.WebSocketGenerator \
      -G websocket_config.json \
      --probes encoding

Testing with Public Echo Server:

garak --model_type websocket.WebSocketGenerator \
      --generator_options '{"websocket": {"WebSocketGenerator": {"uri": "wss://echo.websocket.org", "response_after_typing": false}}}' \
      --probes dan --generations 1

SSH Tunnel Support

The generator works seamlessly with SSH tunnels for secure remote testing:

# Establish tunnel
ssh -L 3000:target-host:3000 jump-host -N -f

# Test through tunnel
garak --model_type websocket.WebSocketGenerator \
      --generator_options '{"websocket": {"WebSocketGenerator": {"uri": "ws://localhost:3000"}}}' \
      --probes malwaregen

Typing Indicators

Many chat-based LLMs send typing indicators. Configure response handling:

response_after_typing: true - Wait for typing to complete (default)
response_after_typing: false - Return first substantial response
typing_indicator - String to detect typing status (default “typing”)

This enables proper testing of streaming/real-time LLM services.

WebSocket generator for real-time LLM communication

This module provides WebSocket-based connectivity for testing LLM services that use real-time bidirectional communication protocols.

class WebSocketGenerator(uri=None, config_root=<module 'garak._config' from '/home/docs/checkouts/readthedocs.org/user_builds/garak/checkouts/latest/docs/source/../../garak/_config.py'>)Source 

Bases: Generator

Generator for WebSocket-based LLM services

This generator connects to LLM services that communicate via WebSocket protocol, handling authentication, template-based messaging, and JSON response extraction.

Configuration parameters: - uri: WebSocket URL (ws:// or wss://) - name: Display name for the service - auth_type: Authentication method (none, basic, bearer, custom) - username: Basic authentication username - api_key: API key for bearer token auth or password for basic auth - req_template: String template with $INPUT and $KEY placeholders - req_template_json_object: JSON object template for structured messages - headers: Additional WebSocket headers - response_json: Whether responses are JSON formatted - response_json_field: Field to extract from JSON responses (supports JSONPath) - response_after_typing: Wait for typing indicator completion - typing_indicator: String that indicates typing status - request_timeout: Seconds to wait for response - connection_timeout: Seconds to wait for connection - max_response_length: Maximum response length - verify_ssl: SSL certificate verification

Configurable parameters:

DEFAULT_PARAMS contents:

uri = 'wss://echo.websocket.org'
name = 'WebSocket LLM'
auth_type = 'none'
username = None
api_key = None
conversation_id = None
req_template = '$INPUT'
req_template_json_object = None
headers = {}
response_json = False
response_json_field = 'text'
response_after_typing = True
typing_indicator = 'typing'
request_timeout = 20
connection_timeout = 10
max_response_length = 10000
verify_ssl = True

Default values are listed

See also Configuring garak for how to set these values.

Other attributes:

ENV_VAR = 'WEBSOCKET_API_KEY'

extra_dependency_names = ['websockets']