FEAT: Adding general websocket target class and accompanying unit tests#1351
Open
kmarsh77 wants to merge 1 commit intoAzure:mainfrom
Open
FEAT: Adding general websocket target class and accompanying unit tests#1351kmarsh77 wants to merge 1 commit intoAzure:mainfrom
kmarsh77 wants to merge 1 commit intoAzure:mainfrom
Conversation
rlundeen2
reviewed
Feb 6, 2026
| logger.info("Successfully connected to websocket") | ||
| return websocket | ||
|
|
||
| async def send_message(self, message: str, conversation_id: str) -> None: |
Contributor
There was a problem hiding this comment.
nit: rename to send_message_async
rlundeen2
reviewed
Feb 6, 2026
| self._existing_conversation = existing_convo if existing_convo is not None else {} | ||
| self._websockets_kwargs = websockets_kwargs or {} | ||
|
|
||
| async def connect(self) -> Any: |
rlundeen2
reviewed
Feb 6, 2026
| endpoint (str): the target endpoint | ||
| initialization_strings (List[str]): These are the connection/initialization strings that must be sent after connecting to websocket in order to initiate conversation | ||
| response_parser: (Callable): Function that takes raw websocket message and tries to parse response message; message is discarded if function fails | ||
| message_builder: (Callable): Function that takes prompt and builds the message to send with it |
Contributor
There was a problem hiding this comment.
Rather than a callable for message_builder, I think this would be more extensible to require a MessageStringNormalizer
rlundeen2
reviewed
Feb 6, 2026
| # Listen for responses | ||
| receive_messages = asyncio.create_task(self.receive_messages(conversation_id=conversation_id)) | ||
|
|
||
| result = await asyncio.wait_for(receive_messages, timeout=30.0) # Wait for all responses to be received |
Contributor
There was a problem hiding this comment.
probably should have this as an init param
rlundeen2
reviewed
Feb 6, 2026
|
|
||
| import pytest | ||
| from websockets.exceptions import ConnectionClosed | ||
| from websockets.frames import Close |
Contributor
There was a problem hiding this comment.
One of my biggest concerns about this is it's hard for us to test. Open to ideas here. In an ideal case, if there is some sort of public socket implementation or something we could chat with that's free or cheap and we're not breaking their TOS.
Then if you added a notebook demoing usage, and chatting with it, it'd be included in our integration tests.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.