diff --git a/packages/sdk/src/index.e2e.test.ts b/packages/sdk/src/index.e2e.test.ts index 504684b..65c64c7 100644 --- a/packages/sdk/src/index.e2e.test.ts +++ b/packages/sdk/src/index.e2e.test.ts @@ -180,7 +180,7 @@ describe('deleteDocument', () => { expect(deletedDoc.projectID).toStrictEqual(newDocument.projectID); expect(await client.getDocuments(testProjectId)).not.toContainEqual(newDocument); - }); + }, 20000); // 20 seconds — test makes 4 sequential HTTP calls }); describe('getDocument', () => { diff --git a/packages/sdk/src/index.test.ts b/packages/sdk/src/index.test.ts index 482dfd8..f3341cc 100644 --- a/packages/sdk/src/index.test.ts +++ b/packages/sdk/src/index.test.ts @@ -87,6 +87,53 @@ describe('MermaidChart', () => { }); }); + describe('#diagramChat', () => { + beforeEach(async () => { + await client.setAccessToken('test-access-token'); + }); + + it('should parse stream body into text and documentChatThreadID', async () => { + const streamBody = [ + '0:"Hello, "', + '0:"here is your diagram!"', + '2:[{"documentChatThreadID":"thread-abc-123"}]', + 'e:{"finishReason":"stop"}', + 'd:{"finishReason":"stop"}', + ].join('\n'); + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + vi.spyOn((client as any).axios, 'post').mockResolvedValue({ data: streamBody }); + + const result = await client.diagramChat({ + message: 'Create a flowchart', + documentID: 'doc-123', + }); + + expect(result.text).toBe('Hello, here is your diagram!'); + expect(result.documentChatThreadID).toBe('thread-abc-123'); + expect(result.documentID).toBe('doc-123'); + }); + + it('should throw AICreditsLimitExceededError on 402', async () => { + // Mock the underlying axios call to simulate a 402 response from the API + // so the actual error-mapping logic inside diagramChat() is exercised. + // eslint-disable-next-line @typescript-eslint/no-explicit-any + vi.spyOn((client as any).axios, 'post').mockRejectedValue({ + response: { + status: 402, + data: 'AI credits limit exceeded', + }, + }); + + await expect( + client.diagramChat({ + message: 'Create a flowchart', + documentID: 'doc-123', + }), + ).rejects.toThrow(AICreditsLimitExceededError); + }); + }); + describe('#repairDiagram', () => { beforeEach(async () => { await client.setAccessToken('test-access-token'); diff --git a/packages/sdk/src/index.ts b/packages/sdk/src/index.ts index 670bd99..cceaacf 100644 --- a/packages/sdk/src/index.ts +++ b/packages/sdk/src/index.ts @@ -10,6 +10,8 @@ import { import type { AuthState, AuthorizationData, + DiagramChatRequest, + DiagramChatResponse, Document, InitParams, MCDocument, @@ -24,6 +26,63 @@ import { URLS } from './urls.js'; const defaultBaseURL = 'https://www.mermaid.ai'; // "http://127.0.0.1:5174" const authorizationURLTimeout = 60_000; +/** + * Parses text tokens from a Vercel AI SDK data-stream response body. + * + * The stream format uses line prefixes: + * `0:"text_chunk"\n` – text token (JSON-encoded string) + * `2:[{"documentChatThreadID":"thread-abc-123"}]\n` – data payload (JSON-encoded array) + * `e:{...}\n` – step finish + * `d:{...}\n` – stream done + */ +function parseVercelAIStreamText(rawBody: string): string { + return rawBody + .split('\n') + .filter((line) => line.startsWith('0:')) + .map((line) => { + try { + const value = JSON.parse(line.slice(2)); + return typeof value === 'string' ? value : ''; + } catch { + return ''; + } + }) + .join(''); +} + +/** + * Extracts data payloads from a Vercel AI SDK data-stream response body. + * Returns the first `documentChatThreadID` found in the stream, if any. + */ +function parseVercelAIStreamData(rawBody: string): { documentChatThreadID?: string } { + let documentChatThreadID: string | undefined; + + for (const line of rawBody.split('\n')) { + if (!line.startsWith('2:')) { + continue; + } + try { + const items: unknown[] = JSON.parse(line.slice(2)); + for (const item of items) { + if (item && typeof item === 'object' && 'documentChatThreadID' in item) { + const value = (item as Record).documentChatThreadID; + if (typeof value === 'string') { + documentChatThreadID = value; + break; + } + } + } + } catch { + // ignore malformed lines + } + if (documentChatThreadID) { + break; + } + } + + return { documentChatThreadID }; +} + export class MermaidChart { private clientID: string; #baseURL!: string; @@ -319,4 +378,79 @@ export class MermaidChart { throw error; } } + + /** + * Chat with Mermaid AI about a diagram. + * + * Sends a single user message to the Mermaid AI chat endpoint. The backend + * automatically fetches the full conversation history from the database + * (when `documentChatThreadID` is provided), so callers never need to track + * or resend previous messages. + * + * @param request - The chat request containing the user message and diagram context + * @returns The AI response text and the chat thread ID + * @throws {@link AICreditsLimitExceededError} if AI credits limit is exceeded (HTTP 402) + */ + public async diagramChat(request: DiagramChatRequest): Promise { + const { message, documentID, code = '', documentChatThreadID } = request; + + // Send only the current user message. The backend will prepend the stored + // conversation history when autoFetchHistory is true (see AIChatRequestData). + const messages = [ + { + id: uuid(), + role: 'user' as const, + content: message, + experimental_attachments: [] as [], + }, + ]; + + const requestBody = { + messages, + code, + documentID, + documentChatThreadID, + // parentID null: the backend already handles finding the correct parent + parentID: null, + // Tell the backend to fetch DB history and prepend it before calling the AI. + autoFetchHistory: true, + }; + + try { + // responseType: 'text' buffers the full stream body as a plain string so we + // can parse the Vercel AI SDK data-stream format after the request completes. + const response = await this.axios.post(URLS.rest.openai.chat, requestBody, { + responseType: 'text', + timeout: 120_000, + }); + + const rawBody = response.data; + const text = parseVercelAIStreamText(rawBody); + const { documentChatThreadID: returnedThreadID } = parseVercelAIStreamData(rawBody); + + return { + text, + documentChatThreadID: returnedThreadID ?? documentChatThreadID, + documentID, + }; + } catch (error: unknown) { + if ( + error && + typeof error === 'object' && + 'response' in error && + error.response && + typeof error.response === 'object' && + 'status' in error.response && + (error as { response: { status: number } }).response.status === 402 + ) { + const axiosError = error as { response: { status: number; data?: unknown } }; + throw new AICreditsLimitExceededError( + typeof axiosError.response.data === 'string' + ? axiosError.response.data + : 'AI credits limit exceeded', + ); + } + throw error; + } + } } diff --git a/packages/sdk/src/types.ts b/packages/sdk/src/types.ts index 9eaec13..f8490cc 100644 --- a/packages/sdk/src/types.ts +++ b/packages/sdk/src/types.ts @@ -96,6 +96,49 @@ export interface RepairDiagramRequest { userID?: string; } +/** + * Request parameters for chatting with the Mermaid AI about a diagram. + */ +export interface DiagramChatRequest { + /** The user's chat message / question */ + message: string; + /** + * MermaidChart document ID to associate the chat thread with. + */ + documentID: string; + /** + * Optional Mermaid diagram code to use as context. + * Pass the current diagram code when the user wants to modify or discuss a specific + * diagram. Leave empty (or omit) for general Mermaid questions. + * Defaults to an empty string. + */ + code?: string; + /** + * Existing chat thread ID to continue a conversation. + * Returned from a previous diagramChat() call. + * When provided, the backend automatically fetches the stored conversation + * history from the database so the AI has full context. + */ + documentChatThreadID?: string; +} + +/** + * Response from chatting with the Mermaid AI. + */ +export interface DiagramChatResponse { + /** The full AI-generated response text (may contain Mermaid code blocks) */ + text: string; + /** + * The document ID used for this conversation. Same as the one passed in the request. + */ + documentID: string; + /** + * The chat thread ID created or used for this conversation. + * Pass this back as documentChatThreadID in subsequent calls to continue the thread. + */ + documentChatThreadID?: string; +} + /** * Response from repairing a diagram. * Matches OpenAIGenerationResult from collab. diff --git a/packages/sdk/src/urls.ts b/packages/sdk/src/urls.ts index 0052e7a..a074a18 100644 --- a/packages/sdk/src/urls.ts +++ b/packages/sdk/src/urls.ts @@ -41,6 +41,7 @@ export const URLS = { }, openai: { repair: `/rest-api/openai/repair`, + chat: `/rest-api/openai/chat`, }, }, raw: (document: Pick, theme: 'light' | 'dark') => {