@kapeta/local-cluster-service 0.67.2 → 0.67.4

package/src/storm/PageGenerator.ts

@@ -20,16 +20,22 @@ export interface ImagePrompt {
     content: string;
 }
 
+type InitialPrompt = Omit<UIPagePrompt, 'shell_page'> & { shellType?: 'public' | 'admin' | 'user' };
+
 export class PageQueue extends EventEmitter {
     private readonly queue: PromiseQueue;
     private readonly systemId: string;
+    private readonly systemPrompt: string;
     private readonly references: Map<string, PageGenerator> = new Map();
+    private readonly pages: Map<string, string> = new Map();
+    private readonly images: Map<string, string> = new Map();
     private uiShells: UIShell[] = [];
     private theme = '';
 
-    constructor(systemId: string, concurrency: number = 5) {
+    constructor(systemId: string, systemPrompt: string, concurrency: number = 5) {
         super();
         this.systemId = systemId;
+        this.systemPrompt = systemPrompt;
         this.queue = new PromiseQueue(concurrency);
     }
 
@@ -56,82 +62,145 @@ export class PageQueue extends EventEmitter {
         this.theme = theme;
     }
 
-    public addPrompt(
-        initialPrompt: Omit<UIPagePrompt, 'shell_page'>,
-        conversationId: string = uuid.v4(),
-        overwrite: boolean = false
-    ) {
-        if (!overwrite && this.references.has(initialPrompt.path)) {
+    private hasPrompt(path: string) {
+        if (this.references.has(path)) {
             //console.log('Ignoring duplicate prompt', initialPrompt.path);
-            return Promise.resolve();
+            return true;
         }
 
-        if (!overwrite && hasPageOnDisk(this.systemId, initialPrompt.method, initialPrompt.path)) {
+        if (hasPageOnDisk(this.systemId, 'GET', path)) {
             //console.log('Ignoring prompt with existing page', initialPrompt.path);
+            return true;
+        }
+
+        return false;
+    }
+
+    public addPrompt(initialPrompt: InitialPrompt, conversationId: string = uuid.v4(), overwrite: boolean = false) {
+        if (!overwrite && this.hasPrompt(initialPrompt.path)) {
+            //console.log('Ignoring duplicate prompt', initialPrompt.path);
             return Promise.resolve();
         }
 
         const prompt: UIPagePrompt = {
             ...initialPrompt,
             shell_page: this.uiShells.find((shell) => shell.screens.includes(initialPrompt.name))?.content,
+            prompt: this.wrapPagePrompt(initialPrompt.path, initialPrompt.prompt),
             theme: this.theme,
         };
 
         const generator = new PageGenerator(prompt, conversationId);
         this.references.set(prompt.path, generator);
+        this.pages.set(prompt.path, prompt.description);
 
         return this.addPageGenerator(generator);
     }
+    private getPrefix(): string {
+        let promptPrefix = '';
+        if (this.systemPrompt) {
+            promptPrefix = `For a system with this description: ${this.systemPrompt}\n`;
+        }
+        return promptPrefix;
+    }
 
-    private async addPageGenerator(generator: PageGenerator) {
-        generator.on('event', (event: StormEvent) => this.emit('event', event));
-        generator.on('page_refs', async ({ event, references }) => {
-            const promises = references.map(async (reference) => {
-                if (
-                    reference.url.startsWith('#') ||
-                    reference.url.startsWith('javascript:') ||
-                    reference.url.startsWith('http://') ||
-                    reference.url.startsWith('https://')
-                ) {
+    private wrapPagePrompt(pagePath: string, prompt: string): string {
+        let promptPrefix = this.getPrefix();
+        let promptPostfix = '';
+
+        if (this.pages.size > 0) {
+            promptPostfix = `\nThe following pages are already implemented:\n`;
+            this.pages.forEach((description, path) => {
+                if (pagePath === path) {
                     return;
                 }
+                promptPostfix += `- PAGE: '${path}' -> ${description}.\n`;
+            });
+        }
 
-                switch (reference.type) {
-                    case 'image':
-                        await this.addImagePrompt({
-                            ...reference,
-                            content: event.payload.content,
-                        });
-                        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: '',
-                            theme: this.theme,
-                        });
-                        break;
-                }
+        if (this.images.size > 0) {
+            promptPostfix += `\nThe following images already exist:\n`;
+            this.images.forEach((description, path) => {
+                promptPostfix += `- IMAGE: '${path}' -> ${description}.\n`;
             });
+        }
+
+        return promptPrefix + prompt + promptPostfix;
+    }
 
-            await Promise.allSettled(promises);
-            this.emit('page', event);
+    private async addPageGenerator(generator: PageGenerator) {
+        generator.on('event', (event: StormEvent) => this.emit('event', event));
+        generator.on('page_refs', async ({ event, references, future }) => {
+            try {
+                const initialPrompts: InitialPrompt[] = [];
+                let promises = references.map(async (reference) => {
+                    if (
+                        reference.url.startsWith('#') ||
+                        reference.url.startsWith('javascript:') ||
+                        reference.url.startsWith('http://') ||
+                        reference.url.startsWith('https://')
+                    ) {
+                        return;
+                    }
+
+                    switch (reference.type) {
+                        case 'image':
+                            await this.addImagePrompt({
+                                ...reference,
+                                content: event.payload.content,
+                            });
+                            break;
+                        case 'css':
+                        case 'javascript':
+                            //console.log('Ignoring reference', reference);
+                            break;
+                        case 'html':
+                            //console.log('Adding page generator for', reference);
+                            this.pages.set(reference.url, reference.description);
+
+                            initialPrompts.push({
+                                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: \n### PATH: ${event.payload.path}\n\`\`\`html\n${event.payload.content}\n\`\`\`\n`,
+                                description: reference.description,
+                                filename: '',
+                                theme: this.theme,
+                            });
+                            break;
+                    }
+                });
+
+                await Promise.allSettled(promises);
+                initialPrompts.forEach((prompt) => {
+                    if (!this.hasPrompt(prompt.path)) {
+                        this.emit('page', {
+                            type: 'PAGE',
+                            reason: 'reference',
+                            created: Date.now(),
+                            payload: {
+                                name: prompt.name,
+                                title: prompt.title,
+                                filename: '',
+                                method: 'GET',
+                                path: prompt.path,
+                                prompt: prompt.description,
+                                conversationId: '',
+                                content: '',
+                                description: prompt.description,
+                            },
+                        });
+                    }
+                    this.addPrompt(prompt);
+                });
+
+                this.emit('page', event);
+            } finally {
+                future.resolve();
+            }
         });
         return this.queue.add(() => generator.generate());
     }
@@ -145,10 +214,18 @@ export class PageQueue extends EventEmitter {
     }
 
     private async addImagePrompt(prompt: ImagePrompt) {
+        if (this.images.has(prompt.url)) {
+            //console.log('Ignoring duplicate image prompt', prompt);
+            return;
+        }
+        this.images.set(prompt.url, prompt.description);
+        const prefix = this.getPrefix();
         const result = await stormClient.createImage(
-            `Create an image for the url "${prompt.url}" with this description: ${prompt.description}`.trim()
+            prefix + `Create an image for the url "${prompt.url}" with this description: ${prompt.description}`.trim()
         );
 
+        //console.log('Adding image prompt', prompt);
+
         const futures: FuturePromise<void>[] = [];
 
         result.on('data', async (event: StormEvent) => {
@@ -172,7 +249,7 @@ export class PageQueue extends EventEmitter {
 
 export class PageGenerator extends EventEmitter {
     private readonly conversationId: string;
-    private prompt: UIPagePrompt;
+    public readonly prompt: UIPagePrompt;
 
     constructor(prompt: UIPagePrompt, conversationId: string = uuid.v4()) {
         super();
@@ -183,7 +260,11 @@ export class PageGenerator extends EventEmitter {
     on(event: 'event', listener: (data: StormEvent) => void): this;
     on(
         event: 'page_refs',
-        listener: (data: { event: StormEventPage; references: ReferenceClassification[] }) => void
+        listener: (data: {
+            event: StormEventPage;
+            references: ReferenceClassification[];
+            future: FuturePromise<void>;
+        }) => void
     ): this;
 
     on(event: string, listener: (...args: any[]) => void): this {
@@ -191,16 +272,18 @@ export class PageGenerator extends EventEmitter {
     }
 
     emit(type: 'event', event: StormEvent): boolean;
-    emit(type: 'page_refs', event: { event: StormEventPage; references: ReferenceClassification[] }): boolean;
+    emit(
+        type: 'page_refs',
+        event: { event: StormEventPage; references: ReferenceClassification[]; future: FuturePromise<void> }
+    ): 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>[] = [];
+            const screenStream = await stormClient.createUIPage(this.prompt, this.conversationId);
 
             screenStream.on('data', (event: StormEvent) => {
                 if (event.type === 'PAGE') {
@@ -210,10 +293,14 @@ export class PageGenerator extends EventEmitter {
                         (async () => {
                             const references = await this.resolveReferences(event.payload.content);
                             //console.log('Resolved references for page', references, event.payload);
+                            const future = createFuture();
                             this.emit('page_refs', {
                                 event,
                                 references,
+                                future,
                             });
+
+                            await future.promise;
                         })()
                     );
                     return;
@@ -222,11 +309,9 @@ export class PageGenerator extends EventEmitter {
                 this.emit('event', event);
             });
 
-            screenStream.on('end', () => {
-                Promise.allSettled(promises).finally(resolve);
-            });
-
             await screenStream.waitForDone();
+            await Promise.all(promises);
+            resolve();
         });
     }
 

package/src/storm/routes.ts

@@ -17,7 +17,6 @@ import {
     ConversationIdHeader,
     stormClient,
     UIPagePrompt,
-    UIPageEditPrompt,
     UIPageEditRequest,
     BasePromptRequest,
     UIPageVoteRequest,
@@ -35,17 +34,11 @@ import {
 import { StormCodegen } from './codegen';
 import { assetManager } from '../assetManager';
 import uuid from 'node-uuid';
-import {
-    readPageFromDisk,
-    readPageFromDiskAsString,
-    SystemIdHeader,
-    writeAssetToDisk,
-    writeImageToDisk,
-    writePageToDisk,
-} from './page-utils';
+import { readPageFromDisk, SystemIdHeader, writeAssetToDisk, writeImageToDisk, writePageToDisk } from './page-utils';
 import { UIServer } from './UIServer';
 import { randomUUID } from 'crypto';
 import { ImagePrompt, PageQueue } from './PageGenerator';
+import { createFuture } from './PromiseQueue';
 
 const UI_SERVERS: { [key: string]: UIServer } = {};
 const router = Router();
@@ -74,7 +67,7 @@ function convertPageEvent(screenData: StormEvent, innerConversationId: string, m
                 description: screenData.payload.description,
                 prompt: screenData.payload.prompt,
                 path: screenData.payload.path,
-                url: server ? server.resolveUrl(screenData) : '',
+                url: server && screenData.payload.content ? server.resolveUrl(screenData) : '',
                 method: screenData.payload.method,
                 conversationId: innerConversationId,
             },
@@ -102,7 +95,7 @@ router.post('/ui/screen', async (req: KapetaBodyRequest, res: Response) => {
 
         const parentConversationId = systemId ?? '';
 
-        const queue = new PageQueue(parentConversationId, 5);
+        const queue = new PageQueue(parentConversationId, '', 5);
         onRequestAborted(req, res, () => {
             queue.cancel();
         });
@@ -138,6 +131,7 @@ router.post('/ui/screen', async (req: KapetaBodyRequest, res: Response) => {
         sendDone(res);
     } catch (err: any) {
         sendError(err, res);
+    } finally {
         if (!res.closed) {
             res.end();
         }
@@ -180,11 +174,17 @@ router.post('/:handle/ui/iterative', async (req: KapetaBodyRequest, res: Respons
         const promises: { [key: string]: Promise<void> } = {};
         const pageEventPromises: Promise<void>[] = [];
         const systemId = landingPagesStream.getConversationId();
-        const pageQueue = new PageQueue(systemId, 5);
+        const systemPrompt = createFuture<string>();
+        if (aiRequest.skipImprovement) {
+            systemPrompt.resolve(aiRequest.prompt);
+        }
 
         landingPagesStream.on('data', async (data: StormEvent) => {
             try {
                 sendEvent(res, data);
+                if (data.type === 'PROMPT_IMPROVE') {
+                    systemPrompt.resolve(data.payload.prompt);
+                }
                 if (data.type !== 'LANDING_PAGE') {
                     return;
                 }
@@ -218,7 +218,13 @@ router.post('/:handle/ui/iterative', async (req: KapetaBodyRequest, res: Respons
 
         UI_SERVERS[systemId] = new UIServer(systemId);
         await UI_SERVERS[systemId].start();
+        waitForStormStream(landingPagesStream).then(() => {
+            if (!systemPrompt.isResolved()) {
+                systemPrompt.resolve(aiRequest.prompt);
+            }
+        });
 
+        const pageQueue = new PageQueue(systemId, await systemPrompt.promise, 5);
         onRequestAborted(req, res, () => {
             pageQueue.cancel();
         });
@@ -281,9 +287,14 @@ router.post('/:handle/ui', async (req: KapetaBodyRequest, res: Response) => {
 
         const uniqueUserJourneyScreens: Record<string, UserJourneyScreen> = {};
 
+        let systemPrompt = aiRequest.prompt;
+
         userJourneysStream.on('data', (data: StormEvent) => {
             try {
                 sendEvent(res, data);
+                if (data.type === 'PROMPT_IMPROVE') {
+                    systemPrompt = data.payload.prompt;
+                }
                 if (data.type !== 'USER_JOURNEY') {
                     return;
                 }
@@ -316,7 +327,6 @@ router.post('/:handle/ui', async (req: KapetaBodyRequest, res: Response) => {
             });
 
             themeStream.on('data', (evt) => {
-                sendEvent(res, evt);
                 if (evt.type === 'FILE_DONE') {
                     theme = evt.payload.content;
                     writeAssetToDisk(outerConversationId, evt).catch((err) => {
@@ -354,7 +364,7 @@ router.post('/:handle/ui', async (req: KapetaBodyRequest, res: Response) => {
         // Get the UI shells
         const shellsStream = await stormClient.createUIShells(
             {
-                theme: theme ? `// filename: theme.css\n${theme}` : undefined,
+                theme: theme || undefined,
                 pages: Object.values(uniqueUserJourneyScreens).map((screen) => ({
                     name: screen.name,
                     title: screen.title,
@@ -371,7 +381,7 @@ router.post('/:handle/ui', async (req: KapetaBodyRequest, res: Response) => {
             shellsStream.abort();
         });
 
-        const queue = new PageQueue(outerConversationId, 5);
+        const queue = new PageQueue(outerConversationId, systemPrompt, 5);
         queue.setUiTheme(theme);
 
         shellsStream.on('data', (data: StormEvent) => {
@@ -410,8 +420,6 @@ router.post('/:handle/ui', async (req: KapetaBodyRequest, res: Response) => {
             created: Date.now(),
         });
 
-        // Get the pages (5 at a time)
-
         const pagePromises: Promise<void>[] = [];
         onRequestAborted(req, res, () => {
             queue.cancel();
@@ -454,17 +462,18 @@ router.post('/:handle/ui', async (req: KapetaBodyRequest, res: Response) => {
             );
         }
 
-        await queue.wait();
-        await Promise.allSettled(pagePromises);
-        await Promise.allSettled(pageEventPromises);
-
         if (userJourneysStream.isAborted()) {
             return;
         }
 
+        await queue.wait();
+        await Promise.allSettled(pagePromises);
+        await Promise.allSettled(pageEventPromises);
+
         sendDone(res);
     } catch (err) {
         sendError(err as Error, res);
+    } finally {
         if (!res.closed) {
             res.end();
         }
@@ -478,7 +487,8 @@ router.post('/ui/edit', async (req: KapetaBodyRequest, res: Response) => {
 
         const aiRequest: StormContextRequest<UIPageEditRequest> = JSON.parse(req.stringBody ?? '{}');
         const storagePrefix = systemId ? systemId + '_' : 'mock_';
-        const queue = new PageQueue(storagePrefix, 5);
+
+        const queue = new PageQueue(systemId!, '', 5);
 
         onRequestAborted(req, res, () => {
             queue.cancel();
@@ -498,8 +508,15 @@ router.post('/ui/edit', async (req: KapetaBodyRequest, res: Response) => {
             }
         });
 
+        const pages = aiRequest.prompt.pages.filter((page) => page.conversationId);
+        if (pages.length === 0) {
+            console.log('No pages to update', aiRequest.prompt.pages);
+            sendDone(res);
+            return;
+        }
+
         await Promise.allSettled(
-            aiRequest.prompt.pages.map((page) => {
+            pages.map((page) => {
                 if (page.conversationId) {
                     return queue.addPrompt(
                         {
@@ -525,6 +542,7 @@ router.post('/ui/edit', async (req: KapetaBodyRequest, res: Response) => {
         sendDone(res);
     } catch (err: any) {
         sendError(err as Error, res);
+    } finally {
         if (!res.closed) {
             res.end();
         }
@@ -535,14 +553,23 @@ router.post('/ui/vote', async (req: KapetaBodyRequest, res: Response) => {
     const conversationId = (req.headers[ConversationIdHeader.toLowerCase()] as string | undefined) || '';
     const aiRequest: UIPageVoteRequest = JSON.parse(req.stringBody ?? '{}');
     const { topic, vote, mainConversationId } = aiRequest;
-    return stormClient.voteUIPage(topic, conversationId, vote, mainConversationId);
+    try {
+        await stormClient.voteUIPage(topic, conversationId, vote, mainConversationId);
+    } catch (e: any) {
+        res.status(500).send({ error: e.message });
+    }
 });
 
 router.post('/ui/get-vote', async (req: KapetaBodyRequest, res: Response) => {
     const conversationId = (req.headers[ConversationIdHeader.toLowerCase()] as string | undefined) || '';
     const aiRequest: UIPageGetVoteRequest = JSON.parse(req.stringBody ?? '{}');
     const { topic, mainConversationId } = aiRequest;
-    return stormClient.getVoteUIPage(topic, conversationId, mainConversationId);
+    try {
+        const vote = await stormClient.getVoteUIPage(topic, conversationId, mainConversationId);
+        res.send({ vote });
+    } catch (e: any) {
+        res.status(500).send({ error: e.message });
+    }
 });
 
 router.post('/:handle/all', async (req: KapetaBodyRequest, res: Response) => {
@@ -689,7 +716,6 @@ router.post('/block/create', async (req: KapetaBodyRequest, res: Response) => {
 
 router.post('/block/codegen', async (req: KapetaBodyRequest, res: Response) => {
     const body: StormCodegenRequest = JSON.parse(req.stringBody ?? '{}');
-    console.log('Codegen request', body);
     const conversationId = req.headers[ConversationIdHeader.toLowerCase()] as string | undefined;
     try {
         const stormCodegen = new StormCodegen(conversationId ?? '', body.prompt, [body.block], body.events || []);
@@ -791,11 +817,14 @@ function onRequestAborted(req: KapetaBodyRequest, res: Response, onAborted: () =
 }
 
 async function sendPageEvent(mainConversationId: string, data: StormEventPage, res: Response) {
-    try {
-        await writePageToDisk(mainConversationId, data);
-    } catch (err) {
-        console.error('Failed to write page to disk', err);
+    if (data.payload.content) {
+        try {
+            await writePageToDisk(mainConversationId, data);
+        } catch (err) {
+            console.error('Failed to write page to disk', err);
+        }
     }
+
     sendEvent(res, convertPageEvent(data, data.payload.conversationId, mainConversationId));
 }
 

package/src/storm/stormClient.ts

@@ -43,7 +43,7 @@ export interface UIPagePrompt {
     description: string;
     storage_prefix: string;
     shell_page?: string;
-    // contents of theme.css
+    // contents of theme.md
     theme?: string;
 }