Allow to interrupt/stop streaming

[krassowski]

Problem

Streaming in chat is great! It allows me to start reading the response before it is fully generated.

However, when I see that the model goes in the wrong way I would like to stop it and ask again.

Proposed Solution

Add a "stop generating" button.

Even though not every provider will support cancelling generation, even having the UI stop moving due to a flush of new text would be an improvement for the UX.

Additional context

Most other chat interfaces have a stop button. For example ChatGPT turns its send button into a stop button while generating.

As does GitHub Copilot:

While the might technically belong to an individual message (since multiple messages could be streamed at the same time), that would likely be a bad UX as streaming would scroll the window moving the button and user would need to chase it.

This could be implemented by extending the SendButton implementation:

jupyter-ai/packages/jupyter-ai/src/components/chat-input/send-button.tsx

Lines 14 to 22 in e45eff5

	export type SendButtonProps = {
	onSend: (selection?: AiService.Selection) => unknown;
	sendWithShiftEnter: boolean;
	currSlashCommand: string | null;
	inputExists: boolean;
	activeCellHasError: boolean;
	};
	export function SendButton(props: SendButtonProps): JSX.Element {

To short-circuit on client-side we can add a breaker condition in:

jupyter-ai/packages/jupyter-ai/src/chat_handler.ts

Lines 145 to 164 in e45eff5

	case 'agent-stream-chunk': {
	const target = newMessage.id;
	const streamMessage = this._messages.find<AiService.AgentStreamMessage>(
	(m): m is AiService.AgentStreamMessage =>
	m.type === 'agent-stream' && m.id === target
	);
	if (!streamMessage) {
	console.error(
	`Received stream chunk with ID ${target}, but no agent-stream message with that ID exists. ` +
	'Ignoring this stream chunk.'
	);
	break;
	}
	streamMessage.body += newMessage.content;
	if (newMessage.stream_complete) {
	streamMessage.complete = true;
	}
	this._messages = [...this._messages];
	break;

It could be a property set on the message. We could then send a request to the server to stop generating messages. We could add a conditional breaker in:

jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/default.py

Lines 104 to 124 in e45eff5

	async for chunk in self.llm_chain.astream(
	{"input": message.body},
	config={"configurable": {"session_id": "static_session"}},
	):
	if not received_first_chunk:
	# when receiving the first chunk, close the pending message and
	# start the stream.
	self.close_pending(pending_message)
	stream_id = self._start_stream(human_msg=message)
	received_first_chunk = True
	if isinstance(chunk, AIMessageChunk):
	self._send_stream_chunk(stream_id, chunk.content)
	elif isinstance(chunk, str):
	self._send_stream_chunk(stream_id, chunk)
	else:
	self.log.error(f"Unrecognized type of chunk yielded: {type(chunk)}")
	break
	# complete stream after all chunks have been streamed
	self._send_stream_chunk(stream_id, "", complete=True)

I am not sure if there is a well-supported way to propagate the interruption to the actual models in langchain. There are some issues, but it does not appear to be documented or officially supported:

• DOC: There is no documentation about how to interrupt the streaming generation once the model started generation. langchain-ai/langchain#11959 suggests raising an exception is suggested
• Implement a way to abort / cancel request langchain-ai/langchain#5340 suggests cancelling the asyncio task