@kapeta/local-cluster-service 0.44.0 → 0.46.0

package/src/storm/codegen.ts

@@ -0,0 +1,266 @@
+/**
+ * Copyright 2023 Kapeta Inc.
+ * SPDX-License-Identifier: BUSL-1.1
+ */
+
+import { Definition } from '@kapeta/local-cluster-config';
+import { AIFileTypes, BlockCodeGenerator, GeneratedFile, GeneratedResult } from '@kapeta/codegen';
+import { BlockDefinition } from '@kapeta/schemas';
+import { codeGeneratorManager } from '../codeGeneratorManager';
+import { STORM_ID, stormClient } from './stormClient';
+import { ScreenTemplate, StormEvent, StormEventFile, StormEventScreen } from './events';
+import { BlockDefinitionInfo } from './event-parser';
+import { ConversationItem, StormFileImplementationPrompt, StormFileInfo, StormStream } from './stream';
+import { KapetaURI } from '@kapeta/nodejs-utils';
+
+type ImplementationGenerator = (
+    prompt: StormFileImplementationPrompt,
+    history?: ConversationItem[]
+) => Promise<StormStream>;
+
+export class StormCodegen {
+    private readonly userPrompt: string;
+    private readonly blocks: BlockDefinitionInfo[];
+    private readonly out = new StormStream();
+    private readonly events: StormEvent[];
+
+    constructor(userPrompt: string, blocks: BlockDefinitionInfo[], events: StormEvent[]) {
+        this.userPrompt = userPrompt;
+        this.blocks = blocks;
+        this.events = events;
+    }
+
+    public async process() {
+        for (const block of this.blocks) {
+            await this.processBlockCode(block);
+        }
+
+        this.out.end();
+    }
+
+    public getStream() {
+        return this.out;
+    }
+
+    private handleTemplateFileOutput(blockUri: KapetaURI, template: StormFileInfo, data: StormEvent) {
+        switch (data.type) {
+            case 'FILE':
+                template.filename = data.payload.filename;
+                template.content = data.payload.content;
+                return this.handleFileOutput(blockUri, data);
+        }
+    }
+
+    private handleUiOutput(blockUri: KapetaURI, data: StormEvent) {
+        switch (data.type) {
+            case 'SCREEN':
+                this.out.emit('data', {
+                    type: 'SCREEN',
+                    reason: data.reason,
+                    created: Date.now(),
+                    payload: {
+                        ...data.payload,
+                        blockName: blockUri.toNormalizedString(),
+                    },
+                });
+            case 'FILE':
+                return this.handleFileOutput(blockUri, data);
+        }
+    }
+
+    private handleFileOutput(blockUri: KapetaURI, data: StormEvent) {
+        switch (data.type) {
+            case 'FILE':
+                this.emitFile(blockUri, data.payload.filename, data.payload.content, data.reason);
+                return {
+                    type: 'FILE',
+                    created: Date.now(),
+                    payload: {
+                        filename: data.payload.filename,
+                        content: data.payload.content,
+                    },
+                } as StormEventFile;
+        }
+    }
+
+    /**
+     * Generates the code for a block and sends it to the AI
+     */
+    private async processBlockCode(block: BlockDefinitionInfo) {
+        // Generate the code for the block using the standard codegen templates
+        const generatedResult = await this.generateBlock(block.content);
+        if (!generatedResult) {
+            return;
+        }
+
+        const allFiles = this.toStormFiles(generatedResult);
+
+        // Send all the non-ai files to the stream
+        this.emitFiles(block.uri, allFiles);
+
+        const relevantFiles: StormFileInfo[] = allFiles.filter(
+            (file) => file.type !== AIFileTypes.IGNORE && file.type !== AIFileTypes.WEB_SCREEN
+        );
+        const uiTemplates: StormFileInfo[] = allFiles.filter((file) => file.type === AIFileTypes.WEB_SCREEN);
+        if (uiTemplates.length > 0) {
+            const uiStream = await stormClient.createUIImplementation({
+                events: this.events,
+                templates: uiTemplates,
+                context: relevantFiles,
+                blockName: block.aiName,
+                prompt: this.userPrompt,
+            });
+
+            uiStream.on('data', (evt) => {
+                this.handleUiOutput(block.uri, evt);
+            });
+
+            await uiStream.waitForDone();
+        }
+
+        // Gather the context files for implementation. These will be all be passed to the AI
+        const contextFiles: StormFileInfo[] = relevantFiles.filter(
+            (file) => ![AIFileTypes.SERVICE, AIFileTypes.WEB_SCREEN].includes(file.type)
+        );
+
+        // Send the service and UI templates to the AI. These will be send one-by-one in addition to the context files
+        const serviceFiles: StormFileInfo[] = allFiles.filter((file) => file.type === AIFileTypes.SERVICE);
+        if (serviceFiles.length > 0) {
+            await this.processTemplates(
+                block.uri,
+                stormClient.createServiceImplementation.bind(stormClient),
+                serviceFiles,
+                contextFiles
+            );
+        }
+    }
+
+    /**
+     * Emits the text-based files to the stream
+     */
+    private emitFiles(uri: KapetaURI, files: StormFileInfo[]) {
+        files.forEach((file) => {
+            if (!file.content || typeof file.content !== 'string') {
+                return;
+            }
+
+            if (file.type === AIFileTypes.SERVICE) {
+                // Don't send the service files to the stream yet
+                // They will need to be implemented by the AI
+                return;
+            }
+
+            if (file.type === AIFileTypes.WEB_SCREEN) {
+                // Don't send the web screen files to the stream yet
+                // They will need to be implemented by the AI
+                return;
+            }
+
+            this.emitFile(uri, file.filename, file.content);
+        });
+    }
+
+    private emitFile(uri: KapetaURI, filename: string, content: string, reason: string = 'File generated') {
+        this.out.emit('data', {
+            type: 'FILE',
+            reason,
+            created: Date.now(),
+            payload: {
+                filename: filename,
+                content: content,
+                blockRef: uri.toNormalizedString(),
+            },
+        } satisfies StormEventFile);
+    }
+
+    /**
+     * Sends the template to the AI and processes the response
+     */
+    private async processTemplates(
+        blockUri: KapetaURI,
+        generator: ImplementationGenerator,
+        templates: StormFileInfo[],
+        contextFiles: StormFileInfo[]
+    ) {
+        const promises = templates.map(async (templateFile) => {
+            const stream = await generator({
+                context: contextFiles,
+                template: templateFile,
+                prompt: this.userPrompt,
+            });
+
+            const files: StormEventFile[] = [];
+
+            stream.on('data', (evt) => {
+                const file = this.handleTemplateFileOutput(blockUri, templateFile, evt);
+                if (file) {
+                    files.push(file);
+                }
+            });
+
+            await stream.waitForDone();
+            return files;
+        });
+
+        const fileChunks = await Promise.all(promises);
+
+        return fileChunks.flat();
+    }
+
+    /**
+     * Converts the generated files to a format that can be sent to the AI
+     */
+    private toStormFiles(generatedResult: GeneratedResult) {
+        const allFiles: StormFileInfo[] = generatedResult.files.map((file: GeneratedFile) => {
+            if (!file.content) {
+                return {
+                    ...file,
+                    type: AIFileTypes.IGNORE,
+                };
+            }
+            if (typeof file.content !== 'string') {
+                return {
+                    ...file,
+                    type: AIFileTypes.IGNORE,
+                };
+            }
+
+            const rx = /\/\/AI-TYPE:([a-z0-9- _]+)\n/gi;
+
+            const match = rx.exec(file.content);
+            if (!match) {
+                return {
+                    ...file,
+                    type: AIFileTypes.IGNORE,
+                };
+            }
+            const type = match[1].trim() as AIFileTypes;
+            file.content = file.content.replace(rx, '');
+
+            return {
+                ...file,
+                type,
+            };
+        });
+        return allFiles;
+    }
+
+    /**
+     * Generates the code using codegen for a given block.
+     */
+    private async generateBlock(yamlContent: Definition) {
+        if (!yamlContent.spec.target?.kind) {
+            //Not all block types have targets
+            return;
+        }
+
+        if (!(await codeGeneratorManager.ensureTarget(yamlContent.spec.target?.kind))) {
+            return;
+        }
+
+        const codeGenerator = new BlockCodeGenerator(yamlContent as BlockDefinition);
+        codeGenerator.withOption('AIContext', STORM_ID);
+
+        return codeGenerator.generate();
+    }
+}