Inline Code Complete with Jupyter-AI Not Working

[passivebook]

Using GPT-4o-mini

The same model works in chat. But for inline code completion I get nothing.

Using the latest version of jupyter-ai and jupyterlabs

Here is a loom video of the problem: https://www.loom.com/share/207ffccdcabd4abe97d81ac83778f15f?sid=89130e09-5294-4d2b-ac9c-24f3309c92d8

[krassowski]

Can you share the logs from Jupyter server (user pod/container if this is running in JupyterHub) and from browser dev tools console from the time when you attempt to get completions?

[passivebook]

Had to reload the entire server. Now it works.

[passivebook]

Inline code completion is a hit or miss. It stopped working again today. I am using JupyterLabs Server in Ubuntu 22.04. Can you please share from where I can get the logs?

[krassowski]

Can you please share from where I can get the logs?

1. See https://webmasters.stackexchange.com/a/77337 for how to access the JavaScript console
2. jupyter-server logs show up in the terminal when you launch JupyterLab

Unless you run JupyterLab Desktop, then see the template here: https://github.com/jupyterlab/jupyterlab-desktop/issues/new?assignees=&labels=bug&projects=&template=bug_report.md

[passivebook]

Are there specific instances where the inline code completion stops firing by design? Like when Kernal is running or if there are multiple notebooks open etc?

[albertmichaelj]

I am having the exact same issue where the inline code completion just does not work after a little bit of time in the notebook. If I reload the browser window, then it continues working again for some period of time. There is nothing in the logs for the server, but if I look in the Javascript console, I get this message when I try to manually invoke the inline completer:

148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1 WebSocket is already in CLOSING or CLOSED state.
(anonymous) @ 148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1
sendMessage @ 148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1
fetch @ 148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1
l @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
fetchInline @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
_makeInlineRequest @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
invokeInline @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
invoke @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
execute @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
execute @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
_executeKeyBinding @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
await in _executeKeyBinding (async)
processKeydownEvent @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
evtKeydown @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
handleEvent @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1 WebSocket is already in CLOSING or CLOSED state.
(anonymous) @ 148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1
sendMessage @ 148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1
fetch @ 148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1
l @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
fetchInline @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
_makeInlineRequest @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
invokeInline @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
invoke @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
execute @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
execute @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
_executeKeyBinding @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
await in _executeKeyBinding (async)
processKeydownEvent @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
evtKeydown @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
handleEvent @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1

When there is a successful inline completion, I get the following message in the server logs, but nothing when the completer does not work.

[I 2024-08-07 14:27:19.736 ServerApp] Inline completion streaming completed in 354 ms.

Note that reloading the browser window does restore functionality for some period of time. This is on a Jupyterhub on k8s instance, so the jupyter lab server is a container running on a kubernetes cluster, if that makes any difference.

[krassowski]

In these cases, does the chat interface continue working or does it fail too?

[albertmichaelj]

It appears that the chat interface is working fine, and it appears to remember previous context just fine. I'm happy to help troubleshoot this in any way possible.

[krassowski]

More logs from the frontend around the time when it starts happening would be useful. Does it include something like "Will try to reconnect in ..."?

For reference, and anyone with ideas, the code for completions websocket handlers is in:

jupyter-ai/packages/jupyter-ai/src/completions/handler.ts

Lines 12 to 149 in 7531f42

	export class CompletionWebsocketHandler implements IDisposable {
	/**
	* The server settings used to make API requests.
	*/
	readonly serverSettings: ServerConnection.ISettings;
	/**
	* Create a new completion handler.
	*/
	constructor(options: AiService.IOptions = {}) {
	this.serverSettings =
	options.serverSettings ?? ServerConnection.makeSettings();
	}
	/**
	* Initializes the WebSocket connection to the completion backend. Promise is
	* resolved when server acknowledges connection and sends the client ID. This
	* must be awaited before calling any other method.
	*/
	public async initialize(): Promise<void> {
	await this._initialize();
	}
	/**
	* Sends a message across the WebSocket. Promise resolves to the message ID
	* when the server sends the same message back, acknowledging receipt.
	*/
	public sendMessage(
	message: AiService.InlineCompletionRequest
	): Promise<AiService.InlineCompletionReply> {
	return new Promise(resolve => {
	this._socket?.send(JSON.stringify(message));
	this._replyForResolver[message.number] = resolve;
	});
	}
	/**
	* Signal emitted when a new chunk of completion is streamed.
	*/
	get streamed(): ISignal<CompletionWebsocketHandler, StreamChunk> {
	return this._streamed;
	}
	/**
	* Whether the completion handler is disposed.
	*/
	get isDisposed(): boolean {
	return this._isDisposed;
	}
	/**
	* Dispose the completion handler.
	*/
	dispose(): void {
	if (this.isDisposed) {
	return;
	}
	this._isDisposed = true;
	// Clean up socket.
	const socket = this._socket;
	if (socket) {
	this._socket = null;
	socket.onopen = () => undefined;
	socket.onerror = () => undefined;
	socket.onmessage = () => undefined;
	socket.onclose = () => undefined;
	socket.close();
	}
	}
	private _onMessage(message: AiService.CompleterMessage): void {
	switch (message.type) {
	case 'connection': {
	this._initialized.resolve();
	break;
	}
	case 'stream': {
	this._streamed.emit(message);
	break;
	}
	default: {
	if (message.reply_to in this._replyForResolver) {
	this._replyForResolver[message.reply_to](message);
	delete this._replyForResolver[message.reply_to];
	} else {
	console.warn('Unhandled message', message);
	}
	break;
	}
	}
	}
	/**
	* Dictionary mapping message IDs to Promise resolvers.
	*/
	private _replyForResolver: Record<
	number,
	(value: AiService.InlineCompletionReply) => void
	> = {};
	private _onClose(e: CloseEvent, reject: any) {
	reject(new Error('Inline completion websocket disconnected'));
	console.error('Inline completion websocket disconnected');
	// only attempt re-connect if there was an abnormal closure
	// WebSocket status codes defined in RFC 6455: https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1
	if (e.code === 1006) {
	const delaySeconds = 1;
	console.info(`Will try to reconnect in ${delaySeconds} s.`);
	setTimeout(async () => await this._initialize(), delaySeconds * 1000);
	}
	}
	private async _initialize(): Promise<void> {
	if (this.isDisposed) {
	return;
	}
	const promise = new PromiseDelegate<void>();
	this._initialized = promise;
	console.log(
	'Creating a new websocket connection for inline completions...'
	);
	const { token, WebSocket, wsUrl } = this.serverSettings;
	const url =
	URLExt.join(wsUrl, SERVICE_URL) +
	(token ? `?token=${encodeURIComponent(token)}` : '');
	const socket = (this._socket = new WebSocket(url));
	socket.onclose = e => this._onClose(e, promise.reject);
	socket.onerror = e => promise.reject(e);
	socket.onmessage = msg => msg.data && this._onMessage(JSON.parse(msg.data));
	}
	private _isDisposed = false;
	private _socket: WebSocket | null = null;
	private _streamed = new Signal<CompletionWebsocketHandler, StreamChunk>(this);
	private _initialized: PromiseDelegate<void> = new PromiseDelegate<void>();
	}

and for chat in:

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

Lines 10 to 256 in 7531f42

	export class ChatHandler implements IDisposable {
	/**
	* The server settings used to make API requests.
	*/
	readonly serverSettings: ServerConnection.ISettings;
	/**
	* ID of the connection. Requires `await initialize()`.
	*/
	id = '';
	/**
	* Create a new chat handler.
	*/
	constructor(options: AiService.IOptions = {}) {
	this.serverSettings =
	options.serverSettings ?? ServerConnection.makeSettings();
	}
	/**
	* Initializes the WebSocket connection to the Chat backend. Promise is
	* resolved when server acknowledges connection and sends the client ID. This
	* must be awaited before calling any other method.
	*/
	public async initialize(): Promise<void> {
	await this._initialize();
	}
	/**
	* Sends a message across the WebSocket. Promise resolves to the message ID
	* when the server sends the same message back, acknowledging receipt.
	*/
	public sendMessage(message: AiService.ChatRequest): Promise<string> {
	return new Promise(resolve => {
	this._socket?.send(JSON.stringify(message));
	this._sendResolverQueue.push(resolve);
	});
	}
	/**
	* Returns a Promise that resolves to the agent's reply, given the message ID
	* of the human message. Should only be called once per message.
	*/
	public replyFor(messageId: string): Promise<AiService.AgentChatMessage> {
	return new Promise(resolve => {
	this._replyForResolverDict[messageId] = resolve;
	});
	}
	public addListener(handler: (message: AiService.Message) => void): void {
	this._listeners.push(handler);
	}
	public removeListener(handler: (message: AiService.Message) => void): void {
	const index = this._listeners.indexOf(handler);
	if (index > -1) {
	this._listeners.splice(index, 1);
	}
	}
	/**
	* Whether the chat handler is disposed.
	*/
	get isDisposed(): boolean {
	return this._isDisposed;
	}
	/**
	* Dispose the chat handler.
	*/
	dispose(): void {
	if (this.isDisposed) {
	return;
	}
	this._isDisposed = true;
	this._listeners = [];
	// Clean up socket.
	const socket = this._socket;
	if (socket) {
	this._socket = null;
	socket.onopen = () => undefined;
	socket.onerror = () => undefined;
	socket.onmessage = () => undefined;
	socket.onclose = () => undefined;
	socket.close();
	}
	}
	get history(): AiService.ChatHistory {
	return {
	messages: this._messages,
	pending_messages: this._pendingMessages
	};
	}
	get historyChanged(): Signal<this, AiService.ChatHistory> {
	return this._historyChanged;
	}
	private _onMessage(newMessage: AiService.Message): void {
	// resolve promise from `sendMessage()`
	if (newMessage.type === 'human' && newMessage.client.id === this.id) {
	this._sendResolverQueue.shift()?.(newMessage.id);
	}
	// resolve promise from `replyFor()` if it exists
	if (
	newMessage.type === 'agent' &&
	newMessage.reply_to in this._replyForResolverDict
	) {
	this._replyForResolverDict[newMessage.reply_to](newMessage);
	delete this._replyForResolverDict[newMessage.reply_to];
	}
	// call listeners in serial
	this._listeners.forEach(listener => listener(newMessage));
	// append message to chat history. this block should always set `_messages`
	// or `_pendingMessages` to a new array instance rather than modifying
	// in-place so consumer React components re-render.
	switch (newMessage.type) {
	case 'connection':
	break;
	case 'clear':
	this._messages = [];
	break;
	case 'pending':
	this._pendingMessages = [...this._pendingMessages, newMessage];
	break;
	case 'close-pending':
	this._pendingMessages = this._pendingMessages.filter(
	p => p.id !== newMessage.id
	);
	break;
	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;
	}
	default:
	// human or agent chat message
	this._messages = [...this._messages, newMessage];
	break;
	}
	// finally, trigger `historyChanged` signal
	this._historyChanged.emit({
	messages: this._messages,
	pending_messages: this._pendingMessages
	});
	}
	/**
	* Queue of Promise resolvers pushed onto by `send()`
	*/
	private _sendResolverQueue: ((value: string) => void)[] = [];
	/**
	* Dictionary mapping message IDs to Promise resolvers, inserted into by
	* `replyFor()`.
	*/
	private _replyForResolverDict: Record<
	string,
	(value: AiService.AgentChatMessage) => void
	> = {};
	private _onClose(e: CloseEvent, reject: any) {
	reject(new Error('Chat UI websocket disconnected'));
	console.error('Chat UI websocket disconnected');
	// only attempt re-connect if there was an abnormal closure
	// WebSocket status codes defined in RFC 6455: https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1
	if (e.code === 1006) {
	const delaySeconds = 1;
	console.info(`Will try to reconnect in ${delaySeconds} s.`);
	setTimeout(async () => await this._initialize(), delaySeconds * 1000);
	}
	}
	private _initialize(): Promise<void> {
	return new Promise<void>((resolve, reject) => {
	if (this.isDisposed) {
	return;
	}
	console.log('Creating a new websocket connection for chat...');
	const { token, WebSocket, wsUrl } = this.serverSettings;
	const url =
	URLExt.join(wsUrl, CHAT_SERVICE_URL) +
	(token ? `?token=${encodeURIComponent(token)}` : '');
	const socket = (this._socket = new WebSocket(url));
	socket.onclose = e => this._onClose(e, reject);
	socket.onerror = e => reject(e);
	socket.onmessage = msg =>
	msg.data && this._onMessage(JSON.parse(msg.data));
	const listenForConnection = (message: AiService.Message) => {
	if (message.type !== 'connection') {
	return;
	}
	this.id = message.client_id;
	// initialize chat history from `ConnectionMessage`
	this._messages = message.history.messages;
	this._pendingMessages = message.history.pending_messages;
	resolve();
	this.removeListener(listenForConnection);
	};
	this.addListener(listenForConnection);
	});
	}
	private _isDisposed = false;
	private _socket: WebSocket | null = null;
	private _listeners: ((msg: any) => void)[] = [];
	/**
	* The list of chat messages
	*/
	private _messages: AiService.ChatMessage[] = [];
	private _pendingMessages: AiService.PendingMessage[] = [];
	/**
	* Signal for when the chat history is changed. Components rendering the chat
	* history should subscribe to this signal and update their state when this
	* signal is triggered.
	*/
	private _historyChanged = new Signal<this, AiService.ChatHistory>(this);
	}

[albertmichaelj]

I was trying to get it to fail while I was actively on the jupyter lab page, but I couldn't. It seems (in my limited testing), that it is most likely to fail when I have navigated to another tab and then come back to Jupyter lab. However, given that, the console in Chrome seems to mix all the messages, and I can't figure out how to limit it to just the single page, so I am not 100% sure what is relevant. When I navigate back to Jupyterlab, I see the following messages at the end of the console log:

148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1 Chat UI websocket disconnected
_onClose @ 148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1
r.onclose @ 148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1
148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1 Will try to reconnect in 1 s.
debug.js:87 (3)(+0001716): Connector_Browser: onActivated for https://github.com/jupyterlab/jupyter-ai/blob/7531f42fc911faf495cc0d8adcd3f9096dafd189/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/aws.py#L16

148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1 Creating a new websocket connection for chat...
debug.js:87 (3)(+0004041): Connector_Browser: onActivated for https://jupyterhub.darden.virginia.edu/gbus-8496-25/user/mja6n/lab/tree/Untitled.ipynb

148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1 WebSocket is already in CLOSING or CLOSED state.
(anonymous) @ 148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1
sendMessage @ 148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1
fetch @ 148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1
l @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
(anonymous) @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
setTimeout (async)
c @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
fetchInline @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
_makeInlineRequest @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
onTextChanged @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
_onSharedModelChanged @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
m @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
l @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
emit @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
S._modelObserver @ 1168.0a95b9c93afe61ceb527.js?v=0a95b9c93afe61ceb527:1
r @ 383.db345dbeef5ef774e50c.js?v=db345dbeef5ef774e50c:1
ji @ 383.db345dbeef5ef774e50c.js?v=db345dbeef5ef774e50c:1
(anonymous) @ 383.db345dbeef5ef774e50c.js?v=db345dbeef5ef774e50c:1
r @ 383.db345dbeef5ef774e50c.js?v=db345dbeef5ef774e50c:1
$o @ 383.db345dbeef5ef774e50c.js?v=db345dbeef5ef774e50c:1
Jo @ 383.db345dbeef5ef774e50c.js?v=db345dbeef5ef774e50c:1
transact @ 383.db345dbeef5ef774e50c.js?v=db345dbeef5ef774e50c:1
update @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
update @ 4986.c819184e09e1e15ef06b.js?v=c819184e09e1e15ef06b:1
updatePlugins @ 4986.c819184e09e1e15ef06b.js?v=c819184e09e1e15ef06b:1
update @ 4986.c819184e09e1e15ef06b.js?v=c819184e09e1e15ef06b:1
dispatchTransactions @ 4986.c819184e09e1e15ef06b.js?v=c819184e09e1e15ef06b:1
dispatch @ 4986.c819184e09e1e15ef06b.js?v=c819184e09e1e15ef06b:1
Fs @ 4986.c819184e09e1e15ef06b.js?v=c819184e09e1e15ef06b:1
flush @ 4986.c819184e09e1e15ef06b.js?v=c819184e09e1e15ef06b:1
(anonymous) @ 4986.c819184e09e1e15ef06b.js?v=c819184e09e1e15ef06b:1
148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1 WebSocket is already in CLOSING or CLOSED state.
(anonymous) @ 148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1
sendMessage @ 148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1
fetch @ 148.aa84a7d464636aa796fe.js?v=aa84a7d464636aa796fe:1
l @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
fetchInline @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
_makeInlineRequest @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
invokeInline @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
invoke @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
execute @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
execute @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
_executeKeyBinding @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
await in _executeKeyBinding (async)
processKeydownEvent @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
evtKeydown @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1
handleEvent @ jlab_core.e5c4c4689689f1651d98.js?v=e5c4c4689689f1651d98:1

This is with me confirming that it has failed by manually trying to invoke inline completion. Does this give more context? Again, I don't want to dump the entire log because it seems to have everything from the browsing session in it.

Chat continues to work just fine.