@kapeta/local-cluster-service 0.61.2 → 0.62.0

package/src/storm/PageGenerator.ts

@@ -0,0 +1,181 @@
+/**
+ * Copyright 2023 Kapeta Inc.
+ * SPDX-License-Identifier: BUSL-1.1
+ */
+
+import uuid from 'node-uuid';
+import { stormClient, UIPagePrompt } from './stormClient';
+import { ReferenceClassification, StormEvent, StormEventPage, StormEventReferenceClassification } from './events';
+import { EventEmitter } from 'node:events';
+import { PromiseQueue } from './PromiseQueue';
+
+export class PageQueue extends EventEmitter {
+    private readonly queue: PromiseQueue;
+    private readonly systemId: string;
+    private readonly references: Map<string, PageGenerator> = new Map();
+
+    constructor(systemId: string, concurrency: number = 5) {
+        super();
+        this.systemId = systemId;
+        this.queue = new PromiseQueue(concurrency);
+    }
+
+    on(event: 'event', listener: (data: StormEvent) => void): this;
+    on(event: 'page', listener: (data: StormEventPage) => void): this;
+
+    on(event: string, listener: (...args: any[]) => void): this {
+        return super.on(event, listener);
+    }
+
+    emit(type: 'event', event: StormEvent): boolean;
+    emit(type: 'page', event: StormEventPage): boolean;
+    emit(eventName: string | symbol, ...args: any[]): boolean {
+        return super.emit(eventName, ...args);
+    }
+
+    public addPrompt(initialPrompt: UIPagePrompt) {
+        if (this.references.has(initialPrompt.path)) {
+            console.log('Ignoring duplicate prompt', initialPrompt.path);
+            return Promise.resolve();
+        }
+        console.log('processing prompt', initialPrompt.path);
+        const generator = new PageGenerator(initialPrompt);
+        this.references.set(initialPrompt.path, generator);
+        return this.addPageGenerator(generator);
+    }
+
+    private async addPageGenerator(generator: PageGenerator) {
+        generator.on('event', (event: StormEvent) => this.emit('event', event));
+        generator.on('page_refs', ({ event, references }) => {
+            this.emit('page', event);
+            references.forEach((reference) => {
+                if (
+                    reference.url.startsWith('#') ||
+                    reference.url.startsWith('javascript:') ||
+                    reference.url.startsWith('http://') ||
+                    reference.url.startsWith('https://')
+                ) {
+                    return;
+                }
+
+                switch (reference.type) {
+                    case 'image':
+                        console.log('Ignoring image reference', reference);
+                        break;
+                    case 'css':
+                    case 'javascript':
+                        //console.log('Ignoring reference', reference);
+                        break;
+                    case 'html':
+                        console.log('Adding page generator for', reference);
+                        const paths = Array.from(this.references.keys());
+                        this.addPrompt({
+                            name: reference.name,
+                            title: reference.title,
+                            path: reference.url,
+                            method: 'GET',
+                            storage_prefix: this.systemId + '_',
+                            prompt:
+                                `Implement a page for ${reference.name} at ${reference.url} with the following description: ${reference.description}.\n` +
+                                `The page was referenced from this page: \`\`\`html\n${event.payload.content}\n\`\`\`\n` +
+                                (paths.length > 0
+                                    ? `\nThese paths are already implemented:\n- ${paths.join('\n - ')}\n\n`
+                                    : ''),
+                            description: reference.description,
+                            filename: '',
+                        });
+                        break;
+                }
+            });
+        });
+        return this.queue.add(() => generator.generate());
+    }
+
+    public cancel() {
+        this.queue.cancel();
+    }
+
+    public wait() {
+        return this.queue.wait();
+    }
+}
+
+export class PageGenerator extends EventEmitter {
+    private readonly conversationId: string;
+    private prompt: UIPagePrompt;
+
+    constructor(prompt: UIPagePrompt, conversationId: string = uuid.v4()) {
+        super();
+        this.conversationId = conversationId;
+        this.prompt = prompt;
+    }
+
+    on(event: 'event', listener: (data: StormEvent) => void): this;
+    on(
+        event: 'page_refs',
+        listener: (data: { event: StormEventPage; references: ReferenceClassification[] }) => void
+    ): this;
+
+    on(event: string, listener: (...args: any[]) => void): this {
+        return super.on(event, listener);
+    }
+
+    emit(type: 'event', event: StormEvent): boolean;
+    emit(type: 'page_refs', event: { event: StormEventPage; references: ReferenceClassification[] }): boolean;
+    emit(eventName: string | symbol, ...args: any[]): boolean {
+        return super.emit(eventName, ...args);
+    }
+
+    public async generate() {
+        return new Promise<void>(async (resolve) => {
+            const screenStream = await stormClient.createUIPage(this.prompt, this.conversationId);
+
+            const promises: Promise<void>[] = [];
+
+            screenStream.on('data', (event: StormEvent) => {
+                if (event.type === 'PAGE') {
+                    event.payload.conversationId = this.conversationId;
+
+                    promises.push(
+                        (async () => {
+                            const references = await this.resolveReferences(event.payload.content);
+                            //console.log('Resolved references for page', references, event.payload);
+                            this.emit('page_refs', {
+                                event,
+                                references,
+                            });
+                        })()
+                    );
+                    return;
+                }
+
+                this.emit('event', event);
+            });
+
+            screenStream.on('end', () => {
+                Promise.allSettled(promises).finally(resolve);
+            });
+
+            await screenStream.waitForDone();
+        });
+    }
+
+    private async resolveReferences(content: string) {
+        const referenceStream = await stormClient.classifyUIReferences(content);
+
+        const references: ReferenceClassification[] = [];
+
+        referenceStream.on('data', (referenceData: StormEvent) => {
+            if (referenceData.type !== 'REF_CLASSIFICATION') {
+                return;
+            }
+
+            //console.log('Processing reference classification', referenceData);
+            references.push(referenceData.payload);
+        });
+
+        await referenceStream.waitForDone();
+
+        return references;
+    }
+}

package/src/storm/events.ts

@@ -406,6 +406,40 @@ export interface StormEventPromptImprove {
     };
 }
 
+export interface LandingPage {
+    name: string;
+    title: string;
+    filename: string;
+    create_prompt: string;
+    path: string;
+    archetype: string;
+    requires_authentication: boolean;
+}
+
+// Event for defining a landing pages
+export interface StormEventLandingPage {
+    type: 'LANDING_PAGE';
+    reason: string;
+    created: number;
+    payload: LandingPage;
+}
+
+export interface ReferenceClassification {
+    name: string;
+    title: string;
+    url: string;
+    description: string;
+    type: 'image' | 'css' | 'javascript' | 'html';
+    source: 'local' | 'cdn' | 'example';
+}
+// Event for reference classification
+export interface StormEventReferenceClassification {
+    type: 'REF_CLASSIFICATION';
+    reason: string;
+    created: number;
+    payload: ReferenceClassification;
+}
+
 export type StormEvent =
     | StormEventCreateBlock
     | StormEventCreateConnection
@@ -435,4 +469,6 @@ export type StormEvent =
     | StormEventUIShell
     | StormEventPage
     | StormEventPageUrl
-    | StormEventPromptImprove;
+    | StormEventPromptImprove
+    | StormEventLandingPage
+    | StormEventReferenceClassification;

package/src/storm/page-utils.ts

@@ -7,18 +7,23 @@ import { Response } from 'express';
 import os from 'node:os';
 import Path from 'path';
 import FS from 'fs-extra';
+import FSExtra from 'fs-extra';
+import { ConversationItem } from './stream';
+import exp from 'node:constants';
 
 export const SystemIdHeader = 'System-Id';
 
+function normalizePath(path: string) {
+    return path
+        .replace(/\?.*$/gi, '')
+        .replace(/:[a-z][a-z_]*\b/gi, '*')
+        .replace(/\{[a-z]+}/gi, '*');
+}
+
 export async function writePageToDisk(systemId: string, event: StormEventPage) {
-    const path = Path.join(
-        os.tmpdir(),
-        'ai-systems',
-        systemId,
-        event.payload.path,
-        event.payload.method.toLowerCase(),
-        'index.html'
-    );
+    const baseDir = getBaseDir(systemId);
+    const filePath = getFilePath(event.payload.method);
+    const path = Path.join(baseDir, normalizePath(event.payload.path), filePath);
     await FS.ensureDir(Path.dirname(path));
     await FS.writeFile(path, event.payload.content);
 
@@ -29,9 +34,58 @@ export async function writePageToDisk(systemId: string, event: StormEventPage) {
     };
 }
 
+function getBaseDir(systemId: string) {
+    return Path.join(os.tmpdir(), 'ai-systems', systemId);
+}
+
+function getFilePath(method: string) {
+    return Path.join(method.toLowerCase(), 'index.html');
+}
+
+export function resolveReadPath(systemId: string, path: string, method: string) {
+    const baseDir = getBaseDir(systemId);
+
+    path = normalizePath(path);
+
+    const filePath = getFilePath(method);
+
+    const fullPath = Path.join(baseDir, path, filePath);
+
+    if (FS.existsSync(fullPath)) {
+        return fullPath;
+    }
+
+    const parts = path.split(/\*/g);
+
+    let currentPath = '';
+
+    for (let part in parts) {
+        const thisPath = Path.join(currentPath, part);
+        const starPath = Path.join(currentPath, '*');
+        const thisPathDir = Path.join(baseDir, thisPath);
+        const starPathDir = Path.join(baseDir, starPath);
+
+        if (FS.existsSync(thisPathDir)) {
+            currentPath = thisPath;
+            continue;
+        }
+
+        if (FS.existsSync(starPathDir)) {
+            currentPath = starPath;
+            continue;
+        }
+
+        console.log('Path not found', thisPathDir, starPathDir);
+        // Path not found
+        return null;
+    }
+
+    return Path.join(baseDir, currentPath, filePath);
+}
+
 export function readPageFromDiskAsString(systemId: string, path: string, method: string) {
-    const filePath = Path.join(os.tmpdir(), 'ai-systems', systemId, path, method.toLowerCase(), 'index.html');
-    if (!FS.existsSync(filePath)) {
+    const filePath = resolveReadPath(systemId, path, method);
+    if (!filePath || !FS.existsSync(filePath)) {
         return null;
     }
 
@@ -39,8 +93,8 @@ export function readPageFromDiskAsString(systemId: string, path: string, method:
 }
 
 export function readPageFromDisk(systemId: string, path: string, method: string, res: Response) {
-    const filePath = Path.join(os.tmpdir(), 'ai-systems', systemId, path, method.toLowerCase(), 'index.html');
-    if (!FS.existsSync(filePath)) {
+    const filePath = resolveReadPath(systemId, path, method);
+    if (!filePath || !FS.existsSync(filePath)) {
         res.status(404).send('Page not found');
         return;
     }
@@ -51,3 +105,25 @@ export function readPageFromDisk(systemId: string, path: string, method: string,
     res.write(content);
     res.end();
 }
+
+export interface Conversation {
+    messages: ConversationItem[];
+    variantId: string;
+    type: 'page';
+    filename: string;
+}
+
+export function readConversationFromFile(filename: string): Conversation[] {
+    if (!FS.existsSync(filename)) {
+        return [];
+    }
+    const content = FS.readFileSync(filename).toString();
+    if (!content.trim()) {
+        return [];
+    }
+    return content.split(/\n/g).map((line) => JSON.parse(line) as Conversation);
+}
+
+export function writeConversationToFile(filename: string, conversations: Conversation[]) {
+    FS.writeFileSync(filename, conversations.map((conversation) => JSON.stringify(conversation)).join('\n'));
+}

package/src/storm/routes.ts

@@ -12,8 +12,17 @@ import { corsHandler } from '../middleware/cors';
 import { stringBody } from '../middleware/stringBody';
 import { KapetaBodyRequest } from '../types';
 import { StormCodegenRequest, StormContextRequest, StormCreateBlockRequest, StormStream } from './stream';
-import { ConversationIdHeader, stormClient, UIPagePrompt, UIPageEditPrompt, UIPageEditRequest } from './stormClient';
+
+import {
+    ConversationIdHeader,
+    stormClient,
+    UIPagePrompt,
+    UIPageEditPrompt,
+    UIPageEditRequest,
+    UIPageSamplePrompt,
+} from './stormClient';
 import { Page, StormEvent, StormEventPage, StormEventPhaseType, UIShell, UserJourneyScreen } from './events';
+
 import {
     createPhaseEndEvent,
     createPhaseStartEvent,
@@ -25,8 +34,17 @@ import { StormCodegen } from './codegen';
 import { assetManager } from '../assetManager';
 import uuid from 'node-uuid';
 import { PromiseQueue } from './PromiseQueue';
-import { readPageFromDisk, readPageFromDiskAsString, SystemIdHeader, writePageToDisk } from './page-utils';
+import {
+    readConversationFromFile,
+    readPageFromDisk,
+    readPageFromDiskAsString,
+    SystemIdHeader,
+    writeConversationToFile,
+    writePageToDisk,
+} from './page-utils';
 import { UIServer } from './UIServer';
+import { PageQueue } from './PageGenerator';
+import FSExtra from 'fs-extra';
 
 const UI_SERVERS: { [key: string]: UIServer } = {};
 const router = Router();
@@ -34,6 +52,8 @@ const router = Router();
 router.use('/', corsHandler);
 router.use('/', stringBody);
 
+const samplesBaseDir = Path.join(__dirname, 'samples');
+
 function convertPageEvent(screenData: StormEvent, innerConversationId: string, mainConversationId: string): StormEvent {
     if (screenData.type === 'PAGE') {
         const server: UIServer | undefined = UI_SERVERS[mainConversationId];
@@ -75,22 +95,25 @@ router.post('/ui/screen', async (req: KapetaBodyRequest, res: Response) => {
         const aiRequest: UIPagePrompt = JSON.parse(req.stringBody ?? '{}');
         aiRequest.storage_prefix = systemId ? systemId + '_' : 'mock_';
 
-        const screenStream = await stormClient.createUIPage(aiRequest, conversationId);
+        res.set('Content-Type', 'application/x-ndjson');
+        res.set('Access-Control-Expose-Headers', ConversationIdHeader);
+        res.set(ConversationIdHeader, conversationId);
+
+        const parentConversationId = systemId ?? '';
 
+        const queue = new PageQueue(parentConversationId, 5);
         onRequestAborted(req, res, () => {
-            screenStream.abort();
+            queue.cancel();
         });
 
-        res.set('Content-Type', 'application/x-ndjson');
-        res.set('Access-Control-Expose-Headers', ConversationIdHeader);
-        res.set(ConversationIdHeader, screenStream.getConversationId());
+        await queue.addPrompt(aiRequest);
 
         const promises: Promise<void>[] = [];
-        screenStream.on('data', (data: StormEvent) => {
+        queue.on('page', (data) => {
             switch (data.type) {
                 case 'PAGE':
                     console.log('Processing page event', data);
-                    data.payload.conversationId = screenStream.getConversationId();
+
                     if (systemId) {
                         promises.push(sendPageEvent(systemId, data, res));
                     }
@@ -99,7 +122,7 @@ router.post('/ui/screen', async (req: KapetaBodyRequest, res: Response) => {
             sendEvent(res, data);
         });
 
-        await waitForStormStream(screenStream);
+        await queue.wait();
         await Promise.allSettled(promises);
 
         sendDone(res);
@@ -125,6 +148,93 @@ router.delete('/:handle/ui', async (req: KapetaBodyRequest, res: Response) => {
     }
 });
 
+router.post('/:handle/ui/iterative', async (req: KapetaBodyRequest, res: Response) => {
+    const handle = req.params.handle as string;
+    try {
+        const conversationId = req.headers[ConversationIdHeader.toLowerCase()] as string | undefined;
+
+        const aiRequest: StormContextRequest = JSON.parse(req.stringBody ?? '{}');
+
+        const landingPagesStream = await stormClient.createUILandingPages(aiRequest.prompt, conversationId);
+
+        onRequestAborted(req, res, () => {
+            landingPagesStream.abort();
+        });
+
+        res.set('Content-Type', 'application/x-ndjson');
+        res.set('Access-Control-Expose-Headers', ConversationIdHeader);
+
+        res.set(ConversationIdHeader, landingPagesStream.getConversationId());
+
+        const promises: { [key: string]: Promise<void> } = {};
+        const pageEventPromises: Promise<void>[] = [];
+        const systemId = landingPagesStream.getConversationId();
+        const pageQueue = new PageQueue(systemId, 5);
+
+        landingPagesStream.on('data', async (data: StormEvent) => {
+            try {
+                sendEvent(res, data);
+                if (data.type !== 'LANDING_PAGE') {
+                    return;
+                }
+
+                if (landingPagesStream.isAborted()) {
+                    return;
+                }
+                const landingPage = data.payload;
+                if (landingPage.name in promises) {
+                    return;
+                }
+
+                // We add the landing pages to the prompt queue.
+                // These will then be analysed - creating further pages as needed
+                promises[landingPage.name] = pageQueue.addPrompt({
+                    prompt: landingPage.create_prompt,
+                    method: 'GET',
+                    path: landingPage.path,
+                    description: landingPage.create_prompt,
+                    name: landingPage.name,
+                    title: landingPage.title,
+                    filename: landingPage.filename,
+                    storage_prefix: systemId + '_',
+                });
+            } catch (e) {
+                console.error('Failed to process event', e);
+            }
+        });
+
+        UI_SERVERS[systemId] = new UIServer(systemId);
+        await UI_SERVERS[systemId].start();
+
+        onRequestAborted(req, res, () => {
+            pageQueue.cancel();
+        });
+
+        pageQueue.on('page', (screenData: StormEventPage) => {
+            pageEventPromises.push(sendPageEvent(landingPagesStream.getConversationId(), screenData, res));
+        });
+
+        pageQueue.on('event', (screenData: StormEvent) => {
+            sendEvent(res, screenData);
+        });
+
+        await waitForStormStream(landingPagesStream);
+        await pageQueue.wait();
+        await Promise.allSettled(pageEventPromises);
+
+        if (landingPagesStream.isAborted()) {
+            return;
+        }
+
+        sendDone(res);
+    } catch (err) {
+        sendError(err as Error, res);
+        if (!res.closed) {
+            res.end();
+        }
+    }
+});
+
 router.post('/:handle/ui', async (req: KapetaBodyRequest, res: Response) => {
     const handle = req.params.handle as string;
     try {
@@ -223,75 +333,40 @@ router.post('/:handle/ui', async (req: KapetaBodyRequest, res: Response) => {
         await UI_SERVERS[outerConversationId].start();
 
         // Get the pages (5 at a time)
-        const queue = new PromiseQueue(5);
-
+        const queue = new PageQueue(outerConversationId, 5);
+        const pagePromises: Promise<void>[] = [];
         onRequestAborted(req, res, () => {
             queue.cancel();
         });
 
+        const pageEventPromises: Promise<void>[] = [];
+        queue.on('page', (screenData: StormEventPage) => {
+            pageEventPromises.push(sendPageEvent(outerConversationId, screenData, res));
+        });
+
+        queue.on('event', (screenData: StormEvent) => {
+            sendEvent(res, screenData);
+        });
+
         for (const screen of Object.values(uniqueUserJourneyScreens)) {
-            void queue.add(
-                () =>
-                    new Promise(async (resolve, reject) => {
-                        try {
-                            const innerConversationId = uuid.v4();
-                            const screenStream = await stormClient.createUIPage(
-                                {
-                                    prompt: screen.requirements,
-                                    method: screen.method,
-                                    path: screen.path,
-                                    description: screen.requirements,
-                                    name: screen.name,
-                                    title: screen.title,
-                                    filename: screen.filename,
-                                    storage_prefix: outerConversationId + '_',
-                                    shell_page: uiShells.find((shell) => shell.screens.includes(screen.name))?.content,
-                                },
-                                innerConversationId
-                            );
-
-                            const promiseList: Promise<void>[] = [];
-                            screenStream.on('data', (screenData: StormEvent) => {
-                                if (screenData.type === 'PAGE') {
-                                    promiseList.push(
-                                        sendPageEvent(
-                                            outerConversationId,
-                                            {
-                                                ...screenData,
-                                                payload: {
-                                                    ...screenData.payload,
-                                                    conversationId: innerConversationId,
-                                                },
-                                            },
-                                            res
-                                        )
-                                    );
-                                } else {
-                                    sendEvent(res, screenData);
-                                }
-                            });
-
-                            screenStream.on('end', async () => {
-                                try {
-                                    await Promise.allSettled(promiseList).finally(() => resolve(true));
-                                } catch (error) {
-                                    console.error('Failed to process screen', error);
-                                }
-                            });
-
-                            screenStream.on('error', (error) => {
-                                console.error('Error on screenStream', error);
-                                screenStream.abort();
-                            });
-                        } catch (e) {
-                            console.error('Failed to process screen', e);
-                            reject(e);
-                        }
-                    })
+            pagePromises.push(
+                queue.addPrompt({
+                    prompt: screen.requirements,
+                    method: screen.method,
+                    path: screen.path,
+                    description: screen.requirements,
+                    name: screen.name,
+                    title: screen.title,
+                    filename: screen.filename,
+                    storage_prefix: outerConversationId + '_',
+                    shell_page: uiShells.find((shell) => shell.screens.includes(screen.name))?.content,
+                })
             );
         }
 
         await queue.wait();
+        await Promise.allSettled(pagePromises);
+        await Promise.allSettled(pageEventPromises);
 
         if (userJourneysStream.isAborted()) {
             return;
@@ -581,6 +656,7 @@ function sendError(err: Error, res: Response) {
         res.status(400).send({ error: err.message });
     }
 }
+
 function waitForStormStream(result: StormStream) {
     return result.waitForDone();
 }