WebSocket Target - PyRIT Documentation

WebsocketTarget connects PyRIT to services that use a custom WebSocket protocol. Supply the service-specific initialization messages, prompt builder, and response parser. The protocol_identifier is a non-secret name for this complete protocol configuration.

This example starts a local PyRIT WebSocket service. It exercises the real WebSocket transport without credentials or an external endpoint.

import json
import uuid

from websockets.asyncio.client import ClientConnection
from websockets.asyncio.server import ServerConnection, serve

from pyrit.models import Message, MessagePiece
from pyrit.prompt_target import WebsocketTarget
from pyrit.setup import IN_MEMORY, initialize_pyrit_async

await initialize_pyrit_async(memory_db_type=IN_MEMORY, load_defaults=False, silent=True)  # type: ignore


async def pyrit_websocket_handler(websocket: ServerConnection) -> None:
    initialization = json.loads(await websocket.recv())
    if initialization != {"type": "initialize", "client": "PyRIT"}:
        await websocket.close(code=1002, reason="Invalid initialization message")
        return

    await websocket.send(json.dumps({"message": "PyRIT WebSocket target ready"}))

    async for raw_message in websocket:
        request = json.loads(raw_message)
        if request["type"] == "restore":
            await websocket.send(json.dumps({"type": "restored"}))
            continue

        await websocket.send(json.dumps({"event": "processing"}))
        await websocket.send(json.dumps({"message": f"PyRIT received: {request['prompt']}"}))


def response_parser(message: str | bytes) -> str | None:
    if isinstance(message, bytes):
        message = message.decode()
    return json.loads(message).get("message")


def message_builder(prompt: str) -> str:
    return json.dumps({"type": "prompt", "prompt": prompt})


async def restore_conversation_async(
    websocket: ClientConnection,
    conversation_history: list[Message],
) -> None:
    history = [
        {
            "role": message.message_pieces[0].role,
            "content": message.get_value(),
        }
        for message in conversation_history
    ]
    await websocket.send(json.dumps({"type": "restore", "history": history}))
    acknowledgement = json.loads(await websocket.recv())
    if acknowledgement != {"type": "restored"}:
        raise ConnectionError("The WebSocket service did not restore the conversation.")

Start the local service on an available loopback port, then configure the target for its protocol.

The restore callback is service-specific. PyRIT calls it when a multi-turn conversation needs a replacement connection. Without this callback, the target fails instead of silently losing history.

server = await serve(pyrit_websocket_handler, "127.0.0.1", 0)  # type: ignore
port = server.sockets[0].getsockname()[1]

target = WebsocketTarget(
    endpoint=f"ws://127.0.0.1:{port}",
    protocol_identifier="local-pyrit-echo-v1",
    initialization_strings=[json.dumps({"type": "initialize", "client": "PyRIT"})],
    response_parser=response_parser,
    message_builder=message_builder,
    conversation_restore_callback=restore_conversation_async,
    discard_initial_messages=1,
)

Send a prompt through the target and close both sides of the connection.

Cleanup is terminal for this target instance. If a connection fails while a prompt is in progress, the target discards that connection and raises the error instead of retrying the prompt.

request = MessagePiece(
    role="user",
    original_value="Hello",
    original_value_data_type="text",
    conversation_id=str(uuid.uuid4()),
).to_message()

try:
    response = await target.send_prompt_async(message=request)  # type: ignore
    print(response[0].get_value())
finally:
    await target.cleanup_target_async()  # type: ignore
    server.close()
    await server.wait_closed()  # type: ignore

PyRIT received: Hello