FEAT: Adding general websocket target class and accompanying unit tests

[kmarsh77]

Description

To address feature request #1037 this commit adds a general WebsocketTarget class for targets using websockets instead of HTTP for communication. As this target class is intended to be generalized to work with all websocket targets, there is some heavy-lifting required by the user which is described below.

After the websocket connection is established, targets typically require several initialization/connection messages for establishing a conversation with the LLM. As such, the WebsocketTarget class requires a list of strings as input containing these messages. After connecting to the provided websocket endpoint, the WebsocketTarget class iterates through and sends the provided initialization messages. These initialization messages can be obtained by connecting to the target normally over a proxy, and extracting the messages from proxy history.

Additionally, there is no standard format for websocket messages. As such, the WebsocketTarget class requires two callable functions as input: response_parser and message_builder. The response parser takes a raw websocket message and extracts the response from LLM; it is expected to fail if message does not contain the actual response, which will allow the discarding of messages sent by server which do not contain the response (e.g. analytics messages). The message builder takes the adversarial prompt as input and returns a formatted websocket message with the injected prompt.

Finally, websocket LLMs typically send 1 or more greeting messages after connection is established. To prevent this message from being interpreted as the response to first adversarial prompt, greeting messages are discarded; the number of initial message to discard is determine by discard_initial_messages argument.

Tests and Documentation

Unit tests have been added for the WebsocketTarget class.

To develop this target class, I used public websites that had chatbots using websockets for communication. I am unsure about the legality of using those chatbots in a working example, especially since it would require the inclusion of server-specific connection strings and message format used by chatbot. Please let me know if there's a better way to approach this.

[Richard Lundeen (rlundeen2)]

Closing due to no movement

[kmarsh77]

Richard Lundeen (@rlundeen2) I had responded to earlier thread asking for your opinion on how to proceed, sorry I should have tagged you.

[Richard Lundeen (rlundeen2)]

Re-opening! kmarsh77

The biggest blocker for me is we need to figure out a way to test it and I'm not sure what that is. One option might be using an existing format like ChatCompletions and calling it using the socket target.

[kmarsh77]

Richard Lundeen (@rlundeen2) so I mentioned above I can make it work with the realtime API (with some minor modifications), that just won't demonstrate the use/purpose of the initialization messages. If you're okay with not being able to test the use of initialization messages, we can go that route. I also had the idea of providing code for a dummy chatbot that communicates using websockets over localhost (and expects initialization strings) but of course this wouldn't be a real-world application.

[kmarsh77]

Richard Lundeen (@rlundeen2) actually, it looks like Lakera's Gandalf Agent Breakers use websockets. I thought about them but just assumed they communicated using HTTP. Considering there is already a target/demo for Gandalf, would it be okay to use the agent breakers as a demo for this websoket target?

[Roman Lutz (romanlutz)]

For a notebook/integration test I think that's a great idea kmarsh77

[kmarsh77]

Okay great, I will get this done in the next week or two when I have some time.

[kmarsh77]

Demo has been added at doc/code/targets/websocket_target.ipynb

[kmarsh77]

Hi Richard Lundeen (@rlundeen2) and Roman Lutz (@romanlutz), just wondering if you've had a chance to look at the demo?