diff --git a/packages/backend/src/managers/playgroundV2Manager.spec.ts b/packages/backend/src/managers/playgroundV2Manager.spec.ts index 2cc121a3c..f9cb7f7be 100644 --- a/packages/backend/src/managers/playgroundV2Manager.spec.ts +++ b/packages/backend/src/managers/playgroundV2Manager.spec.ts @@ -26,6 +26,7 @@ import { Messages } from '@shared/Messages'; import type { ModelInfo } from '@shared/src/models/IModelInfo'; import type { TaskRegistry } from '../registries/TaskRegistry'; import type { Task, TaskState } from '@shared/src/models/ITask'; +import type { ChatMessage, ErrorMessage } from '@shared/src/models/IPlaygroundMessage'; vi.mock('openai', () => ({ default: vi.fn(), @@ -185,7 +186,7 @@ test('valid submit should create IPlaygroundMessage and notify the webview', asy // Wait for assistant message to be completed await vi.waitFor(() => { - expect(manager.getConversations()[0].messages[1].content).toBeDefined(); + expect((manager.getConversations()[0].messages[1] as ChatMessage).content).toBeDefined(); }); const conversations = manager.getConversations(); @@ -272,6 +273,72 @@ test('submit should send options', async () => { }); }); +test('error', async () => { + vi.mocked(inferenceManagerMock.getServers).mockReturnValue([ + { + status: 'running', + health: { + Status: 'healthy', + }, + models: [ + { + id: 'dummyModelId', + file: { + file: 'dummyModelFile', + }, + }, + ], + connection: { + port: 8888, + }, + } as unknown as InferenceServer, + ]); + const createMock = vi.fn().mockRejectedValue('Please reduce the length of the messages or completion.'); + vi.mocked(OpenAI).mockReturnValue({ + chat: { + completions: { + create: createMock, + }, + }, + } as unknown as OpenAI); + + const manager = new PlaygroundV2Manager(webviewMock, inferenceManagerMock, taskRegistryMock, telemetryMock); + await manager.createPlayground('playground 1', { id: 'dummyModelId' } as ModelInfo, 'tracking-1'); + + const date = new Date(2000, 1, 1, 13); + vi.setSystemTime(date); + + const playgrounds = manager.getConversations(); + await manager.submit(playgrounds[0].id, 'dummyUserInput'); + + // Wait for error message + await vi.waitFor(() => { + expect((manager.getConversations()[0].messages[1] as ErrorMessage).error).toBeDefined(); + }); + + const conversations = manager.getConversations(); + + expect(conversations.length).toBe(1); + expect(conversations[0].messages.length).toBe(2); + expect(conversations[0].messages[0]).toStrictEqual({ + content: 'dummyUserInput', + id: expect.anything(), + options: undefined, + role: 'user', + timestamp: expect.any(Number), + }); + expect(conversations[0].messages[1]).toStrictEqual({ + error: 'Please reduce the length of the messages or completion. Note: You should start a new playground.', + id: expect.anything(), + timestamp: expect.any(Number), + }); + + expect(webviewMock.postMessage).toHaveBeenLastCalledWith({ + id: Messages.MSG_CONVERSATIONS_UPDATE, + body: conversations, + }); +}); + test('creating a new playground should send new playground to frontend', async () => { vi.mocked(inferenceManagerMock.getServers).mockReturnValue([]); const manager = new PlaygroundV2Manager(webviewMock, inferenceManagerMock, taskRegistryMock, telemetryMock); @@ -570,7 +637,7 @@ describe('system prompt', () => { // Wait for assistant message to be completed await vi.waitFor(() => { - expect(manager.getConversations()[0].messages[1].content).toBeDefined(); + expect((manager.getConversations()[0].messages[1] as ChatMessage).content).toBeDefined(); }); expect(() => { diff --git a/packages/backend/src/managers/playgroundV2Manager.ts b/packages/backend/src/managers/playgroundV2Manager.ts index d6f770285..7fa8c8911 100644 --- a/packages/backend/src/managers/playgroundV2Manager.ts +++ b/packages/backend/src/managers/playgroundV2Manager.ts @@ -22,8 +22,14 @@ import type { ChatCompletionChunk, ChatCompletionMessageParam } from 'openai/src import type { ModelOptions } from '@shared/src/models/IModelOptions'; import type { Stream } from 'openai/streaming'; import { ConversationRegistry } from '../registries/ConversationRegistry'; -import type { Conversation, PendingChat, SystemPrompt, UserChat } from '@shared/src/models/IPlaygroundMessage'; -import { isSystemPrompt } from '@shared/src/models/IPlaygroundMessage'; +import type { + Conversation, + ErrorMessage, + PendingChat, + SystemPrompt, + UserChat, +} from '@shared/src/models/IPlaygroundMessage'; +import { isChatMessage, isSystemPrompt } from '@shared/src/models/IPlaygroundMessage'; import type { ModelInfo } from '@shared/src/models/IModelInfo'; import { withDefaultConfiguration } from '../utils/inferenceUtils'; import { getRandomString } from '../utils/randomUtils'; @@ -238,12 +244,29 @@ export class PlaygroundV2Manager implements Disposable { .catch((err: unknown) => { telemetry['errorMessage'] = `${String(err)}`; console.error('Something went wrong while creating model response', err); + this.processError(conversation.id, err).catch((err: unknown) => { + console.error('Something went wrong while processing stream', err); + }); }) .finally(() => { this.telemetry.logUsage('playground.submit', telemetry); }); } + private async processError(conversationId: string, error: unknown): Promise { + let errorMessage = String(error); + if (errorMessage.endsWith('Please reduce the length of the messages or completion.')) { + errorMessage += ' Note: You should start a new playground.'; + } + const messageId = this.#conversationRegistry.getUniqueId(); + const start = Date.now(); + this.#conversationRegistry.submit(conversationId, { + id: messageId, + timestamp: start, + error: errorMessage, + } as ErrorMessage); + } + /** * Given a Stream from the OpenAI library update and notify the publisher * @param conversationId @@ -283,13 +306,15 @@ export class PlaygroundV2Manager implements Disposable { const conversation = this.#conversationRegistry.get(conversationId); if (!conversation) throw new Error(`conversation with id ${conversationId} does not exist.`); - return conversation.messages.map( - message => - ({ - name: undefined, - ...message, - }) as ChatCompletionMessageParam, - ); + return conversation.messages + .filter(m => isChatMessage(m)) + .map( + message => + ({ + name: undefined, + ...message, + }) as ChatCompletionMessageParam, + ); } getConversations(): Conversation[] { diff --git a/packages/backend/src/registries/ConversationRegistry.ts b/packages/backend/src/registries/ConversationRegistry.ts index c9d4263de..1206ba6d5 100644 --- a/packages/backend/src/registries/ConversationRegistry.ts +++ b/packages/backend/src/registries/ConversationRegistry.ts @@ -22,6 +22,7 @@ import type { ChatMessage, Choice, Conversation, + Message, PendingChat, } from '@shared/src/models/IPlaygroundMessage'; import type { Disposable, Webview } from '@podman-desktop/api'; @@ -141,7 +142,7 @@ export class ConversationRegistry extends Publisher implements D * @param conversationId * @param message */ - submit(conversationId: string, message: ChatMessage): void { + submit(conversationId: string, message: Message): void { const conversation = this.#conversations.get(conversationId); if (conversation === undefined) throw new Error(`conversation with id ${conversationId} does not exist.`); diff --git a/packages/frontend/src/pages/Playground.svelte b/packages/frontend/src/pages/Playground.svelte index 48c874c13..a054e34c6 100644 --- a/packages/frontend/src/pages/Playground.svelte +++ b/packages/frontend/src/pages/Playground.svelte @@ -1,7 +1,14 @@