aimyapi 0.2.1 → 1.0.2

package/src/index.ts

@@ -3,177 +3,259 @@ import fs from 'fs';
 import { OpenAI } from "openai";
 import path from 'path';
 
-import createSandbox from "./sandbox";
+import createSandbox, { CodeRunResult } from "./sandbox";
 import { rejectOpenPromises } from './sandbox-wrappers';
 
 import { ChatCompletionMessageParam } from 'openai/resources/index.mjs';
 import { QuickJSWASMModule, getQuickJS } from "quickjs-emscripten";
 
-const openai = new OpenAI({
-    apiKey: process.env.OPENAI_API_KEY // This is also the default, can be omitted
-  });
+import { AutoParseableTool } from 'openai/lib/parser.js';
+import { ZodSchema, z } from 'zod';
+import { zodToJsonSchema } from 'zod-to-json-schema';
 
-const CODE_START = `\`\`\`
-import * as ApiDefs from 'api.ts'
+const openai = new OpenAI();
 
-(async() {
-    try {`;
+let QuickJS:QuickJSWASMModule | null = null;
 
-const CODE_END = `    } catch(err) { 
+export interface AIMyAPIOptions {
+    apiObject: object;
+    apiExports: object;
+    apiWhitelist?: string[];      // List of functions to expose to the user. Necessary to pass this yourself if you are extending your api from another class.
+    apiGlobals?: object;
+    apiDefFilePath: string;     // Full path to the api definition file so that it can be loaded into the prompt
+    apiGlobalName?: string;    
+    apiDocsPath?: string;
+    debug?: boolean;
+    model?: string;
+    additionalModelOptions?: object;
+    hideLogsFromAgent?: boolean; // If true, logs from code execution will not be included in the messages sent to the agent. This can be useful to prevent the agent from being overwhelmed with logs or to hide implementation details. Default is false (logs will be included in messages).
+};
+
+export interface AIMyAPIInstance {
+    options:AIMyAPIOptions;
+    generateCode: (queryText:string, userChatHistory:ChatCompletionMessageParam[]) => Promise<GenerateCodeResult | null>;
+    runCode: (code:string) => Promise<CodeRunResult>;
+    checkCode: (code:string) => Promise<CodeRunResult>;
+    // run a single request with no history
+    processRequest: (userQuery:string, context?: object) => Promise<GenerateCodeResult | null>;
+}
+
+export interface GenerateCodeResult {
+    code?: string;
+    comments?: string;
+}
+
+
+const wrapCode = (apiGlobalName:string, apiFileDefPath:string, code:string) => `
+(async () => {
+    try {
+        ${code}
+    } catch(err) { 
         console.error(err); 
     }
 })();        
-\`\`\``;
-
-function indentString(str, numTabs) {
-    const tabs = '\t'.repeat(numTabs);
-    return str.split('\n').map(line => tabs + line).join('\n');
-  }
-
-function replaceCodeBlocks(text) {
-    const codeBlockRegex = /```[\s\S]*?```/g;
-    return text.replace(codeBlockRegex, (match) => {
-        const code = match.slice(3, -3).trim();
-        return `${CODE_START}\n${indentString(code, 2)}\n${CODE_END}`;
-});
+`;
+
+const wrapCodeBlock = (code:string, codeType:string = "") => `
+\`\`\`${codeType}
+${code}
+\`\`\`
+`;
+
+export function zodParseJSON<T>(schema: ZodSchema<T>) {
+  return (input: string): T => schema.parse(JSON.parse(input));
 }
 
-export function createBasePrompt(apiFilePath:string, documentationPath: string): string {
+export function createBasePrompt(apiFilePath:string, apiGlobalName: string, documentationPath?: string, ): string {
     const apiText: string = fs.readFileSync(apiFilePath, 'utf8');
     const documentationText: string = documentationPath ? fs.readFileSync(documentationPath, 'utf8') : '';
 
     //console.time("Loading templates");
-    const createTaskPrompt: string = fs.readFileSync( path.join(__dirname, '../prompts/create-task-prompt.md'), 'utf8').replace("{{DOCUMENTATION}}", replaceCodeBlocks(documentationText)).replace("{{API}}", apiText);   
+    const createTaskPrompt: string = fs.readFileSync( path.join(__dirname, '../prompts/create-task-prompt.md'), 'utf8').replace("{{DOCUMENTATION}}", documentationText).replace("{{API}}", apiText).replace("{{API_GLOBAL_NAME}}", apiGlobalName || "");   
     //console.timeEnd("Loading templates")
     
     return createTaskPrompt;
 }
 
-export interface GenerateCodeResult {
-    code: string;
-    loggableCode: string;
-    rawResponse?: string;
-}
 
-// FIXME: add an options parameter and allow specifying the token limit
-export const generateCode = async function(queryText:string, userChatHistory:ChatCompletionMessageParam[], createTaskPrompt:string, debug:boolean = false, model="gpt-4o-mini"): Promise<GenerateCodeResult> {
-    if(!queryText)
-        return { code: '', loggableCode: '' };
+const runCodeSchema = z.object({
+    code: z.string(),
+});
+
+const checkCodeSchema = z.object({
+    code: z.string(),
+});
+
+const submitCodeSchema = z.object({
+    code: z.string(),
+});
 
-    const prompt = createTaskPrompt.replace("{{QUERY_TEXT}}", queryText);
-    const messages:ChatCompletionMessageParam[] = [
+export const generateCode = async function(instance:AIMyAPIInstance, queryText:string, userChatHistory:ChatCompletionMessageParam[], createTaskPrompt:string, debug:boolean = false, hideLogsFromAgent:boolean = false, model="gpt-5-mini", additionalModelOptions?: object): Promise<GenerateCodeResult | null> {
+    let finalSubmittedCode = "";
+
+    if(debug) {
+        console.log("Generating code for query:", queryText);
+    }
+    // checkcode
+    const codingTools: AutoParseableTool<any, true>[] = [
         {
-            role: "system",
-            content: prompt,
+            type: "function",
+            function: {
+                name: "runCode",
+                // @ts-ignore
+                function: async ({code}: {code: string}) => {   
+                    if(debug) {
+                        console.log("Running code...", code);                 
+                    }
+                    
+                    const res = await instance.runCode(code);
+
+                    if(hideLogsFromAgent) {
+                        res.logs = [];
+                    }
+                    return res;
+                },
+                description: "Run the provided code in a sandboxed environment. The code will have access to the API object. Returns the result of the code execution.",
+                parse: zodParseJSON(runCodeSchema),
+                parameters: zodToJsonSchema(runCodeSchema),            
+            }
         },
-        ...userChatHistory,
         {
-            role: "user",
-            content: `${queryText}. Respond with json`,
-            name: "user"
+            type: "function",
+            function: {
+                name: "checkCode",
+                // @ts-ignore
+                function: async ({code}: {code: string}) => {   
+                    if(debug) {
+                        console.log("Checking code...", code);                 
+                    }
+                    const res = await instance.checkCode(code);
+                    
+                    return res;
+
+                },
+                description: "Check the provided code for errors or issues without running it. Returns the result of the code check.",
+                parse: zodParseJSON(checkCodeSchema),
+                parameters: zodToJsonSchema(checkCodeSchema),
+            }
+        },        
+        {
+            type: "function",
+            function: {
+                name: "submitCode",
+                // @ts-ignore
+                function: async ({code}: {code: string}) => {   
+                    finalSubmittedCode = code;
+                    return {
+                        success: true,
+                    }
+                },
+                description: "Submit the provided code as the final code to run. This will end the code generation process and the code will be run in the sandbox. Use this when you are confident in the code you have generated and want to run it.",
+                parse: zodParseJSON(submitCodeSchema),
+                parameters: zodToJsonSchema(submitCodeSchema),
+            }
         }
     ];
 
     try {
-        const results = await openai.chat.completions.create({
-            model,
-            messages: messages,
-            temperature: 0.0,
-            max_tokens: 2000,
-            response_format: {
-                json_schema: {
-                    name: "codeResponse",
-                    description: "The code generated by the AI",
-                    schema: {
-                        type: 'object',
-                        properties: {
-                            yourCode: {
-                                type: 'string',
-                                description: 'The code generated by the AI',
-                            }
-                        },
-                    },
-                },
-                type: 'json_schema',
-              }
+        
+        const runner = openai.chat.completions
+        .runTools({
+          model,
+          messages: [
+            {
+              role: 'system',
+              content: "You will assist with writing code to perform the user's request.",
+            },
+            ...userChatHistory,
+            {
+                role: 'user',
+                content: createTaskPrompt.replace("{{USER_QUERY}}", queryText)
+            },
+
+          ],
+          tools: codingTools,
+          ...additionalModelOptions,
+
+        })
+        .on('message', async (message: ChatCompletionMessageParam) => {
+            if(debug) {
+                console.log("Message:", JSON.stringify(message, null, 2));
+
+            }
+        })
+        .on('functionToolCall', (functionCall: OpenAI.Chat.Completions.ChatCompletionMessageFunctionToolCall.Function) => {
+            if(debug) {
+                console.log("Function Call:", JSON.stringify(functionCall, null, 2));
             }
-        );
+            
 
-        if(debug)
-            console.log(results.usage);
-        
-        const response = results.choices[0].message.content;
+        })
+        .on('functionToolCallResult', (functionCall) => {
+            if(debug) {
+                console.log("Function Call Result:", JSON.stringify(functionCall, null, 2));
+            }
+        })
+    
+      const finalContent = await runner.finalContent();
+
+      if(debug) {
+        console.log("Final Content:", finalContent);
+      }
+      
+      return {
+        code: finalSubmittedCode,
+        comments: finalContent || undefined,
+      }
 
-        if(debug)
-            console.log(response);
-        
-        try {
-            const json = JSON.parse(response); 
-            const code = json.yourCode.replaceAll('```', '');
-
-            return { code, loggableCode: code };
-        } catch(err) {
-            console.error("Error parsing response", err);
-            return { code: '', loggableCode: '' };
-        }
     } catch(err) {
         console.error(err);
-        return { code: '', loggableCode: '' };
-    } 
-};
-
-let QuickJS:QuickJSWASMModule = null;
-
-export interface AIMyAPIOptions {
-    apiObject: object;
-    apiExports: object;
-    apiWhitelist?: string[];      // List of functions to expose to the user. Necessary to pass this yourself if you are extending your api from another class.
-    apiGlobals?: object;
-    apiDefFilePath: string;     // Full path to the api definition file so that it can be loaded into the prompt
-    apiGlobalName?: string;    
-    apiDocsPath?: string;
-    debug?: boolean;
-    model?: string;
+        return null;
+    }
 };
 
-export interface AIMyAPIInstance {
-    // see xample 2
-    options:AIMyAPIOptions;
-    generateCode: (queryText:string, userChatHistory:ChatCompletionMessageParam[]) => Promise<GenerateCodeResult>;
-    runCode: (task:string) => Promise<void>;
-    // run a single request with no history
-    processRequest: (userQuery:string, context?: object) => Promise<GenerateCodeResult>;
-    
-}
-
 async function createWithAPI(options:AIMyAPIOptions): Promise<AIMyAPIInstance> {
     // extend options with defaults
-    options = {
-        apiGlobalName: "api",
-        apiGlobals: {},
-        apiWhitelist: Object.getOwnPropertyNames(Object.getPrototypeOf(options.apiObject)).filter((f) => f !== "constructor" && !f.startsWith("_")),
-        debug: false, 
-        model: "gpt-4o-mini",
+    const localOptions = {
+        ...{
+            apiGlobalName: "api",
+            apiGlobals: {},
+            apiWhitelist: Object.getOwnPropertyNames(Object.getPrototypeOf(options.apiObject)).filter((f) => f !== "constructor" && !f.startsWith("_")),
+            debug: false, 
+            model: process.env.OPENAI_API_MODEL ||"gpt-5.4-mini",
+            additionalModelOptions: {},
+            hideLogsFromAgent: false,
+        },
         ...options
-    } as AIMyAPIOptions;
+    };
 
-    if(!options.apiObject || !options.apiDefFilePath) {
-        throw new Error("apiObject and apiFilePath are required");
+    if(!localOptions.apiObject || !localOptions.apiDefFilePath) {
+        throw new Error("apiObject and apiDefFilePath are required");
     }
 
-    const {apiObject, apiWhitelist, apiGlobalName, apiExports, apiDefFilePath, apiDocsPath, debug, model} = options;
+    const {apiObject, apiWhitelist, apiGlobalName, apiExports, apiDefFilePath, apiDocsPath, debug, model, additionalModelOptions, hideLogsFromAgent} = localOptions;
 
-    const createTaskPrompt = createBasePrompt( apiDefFilePath, apiDocsPath);
+    const createTaskPrompt = createBasePrompt( apiDefFilePath, apiGlobalName, apiDocsPath);
 
-    return {
-        options,
-        generateCode: async function(queryText:string, userChatHistory:ChatCompletionMessageParam[], currentContext:any = undefined) {
+    const instance = {
+        options: localOptions,
+        generateCode: async function(queryText:string, userChatHistory:ChatCompletionMessageParam[], currentContext:any = undefined) {            
             if(!queryText)
-                return { code: '', loggableCode: '' };
-
-            return await generateCode(queryText, userChatHistory, createTaskPrompt.replace("{{CONTEXT}}", currentContext ? "```" + JSON.stringify(currentContext, null, 2) + "```" : ""), debug, model);
+                return null;
+
+            const formattedContext = wrapCodeBlock(JSON.stringify(currentContext, null, 2), "json");
+
+            return await generateCode(instance, 
+                queryText, 
+                userChatHistory, 
+                createTaskPrompt.replace("{{CONTEXT}}", formattedContext), 
+                debug, 
+                hideLogsFromAgent, 
+                model, 
+                additionalModelOptions
+            );
         },
-        runCode: async function(generatedCode:string) {
+        checkCode: async function (generatedCode:string): Promise<CodeRunResult> {
             if(!QuickJS) {
                 // Create QuickJS
                 if(debug)
@@ -181,18 +263,93 @@ async function createWithAPI(options:AIMyAPIOptions): Promise<AIMyAPIInstance> {
                 QuickJS = await getQuickJS();
             }
 
-            let apiGlobals = {}
+            let apiGlobals: { [key: string]: { value: any } } = {}
 
-            for(const key in options.apiGlobals) {
-                apiGlobals[key] = {
-                    value: options.apiGlobals[key],
+            if(options.apiGlobals) {
+                for(const key in options.apiGlobals) {
+                    apiGlobals[key] = {
+                        value: (options.apiGlobals as Record<string, any>)[key],
+                    }
                 }
             }
-      
+        
+            // Create sandbox           
+            const {vm, checkCode} = await createSandbox(QuickJS, {
+                [apiDefFilePath]: apiExports,
+            }, { 
+                ...apiGlobals,
+                [apiGlobalName]: {
+                    value: apiObject,
+                    whitelist: apiWhitelist
+                }
+            }, options.debug);
+
+            try {
+                if(debug) {
+                    console.log("checking Code:\n", generatedCode);
+                }
+
+                const res = await checkCode(wrapCode(apiGlobalName, apiDefFilePath, generatedCode));
+        
+                if(!res) {
+                    console.error("Unknown failure checking code");
+                    return {
+                        success: false,
+                        logs: [],
+                    };
+                } 
+        
+            
+                if(debug) {
+                    console.log(res.success ? "Code check passed" : "Code check failed");
+                }
+
+                if(res.success) {
+                    return {
+                        success: true,
+                        logs: [{ type: "log", args: ["✔️ Code check passed"], source: "compiler" }, ...res.logs],
+                    }
+                }
+
+                return {
+                    success: false,
+                    logs: [{ type: "error", args: ["❌ Code check failed"], source: "compiler" }, ...res.logs],
+                };
+            } catch(err) {
+                if(debug) {
+                    console.error("Error checking code:", err);
+                }
+            }
+            vm.dispose();
+
+            return {
+                success: false,
+                logs: [],
+            };
+        },    
+        runCode: async function (code:string): Promise<CodeRunResult> {
+            // TODO: fix code duplication with checkCode, maybe combine into a single function that takes a parameter for whether to actually run the code or just check it?
+            if(!QuickJS) {
+                // Create QuickJS
+                if(debug)
+                    console.log("Creating quick.js");
+                QuickJS = await getQuickJS();
+            }
+
+            let apiGlobals: { [key: string]: { value: any } } = {}
+
+            if(options.apiGlobals) {
+                for(const key in options.apiGlobals) {
+                    apiGlobals[key] = {
+                        value: (options.apiGlobals as Record<string, any>)[key],
+                    }
+                }
+            }
+        
             // Create sandbox           
             // NOTE: not sure why I have to replace the backslash here
-            const {vm, runTask, isAsyncProcessRunning} = await createSandbox(QuickJS, {
-                [apiDefFilePath.replaceAll("\\", "")]: apiExports,
+            const {vm, runCode: runTask, isAsyncProcessRunning} = await createSandbox(QuickJS, {
+                [apiDefFilePath]: apiExports,
             }, { 
                 ...apiGlobals,
                 [apiGlobalName]: {
@@ -203,13 +360,17 @@ async function createWithAPI(options:AIMyAPIOptions): Promise<AIMyAPIInstance> {
 
             try {
                 if(debug) {
-                    console.log("Running Code:\n", generatedCode);
+                    console.log("Running Code:\n", code);
                 }
 
-                const res = await runTask(generatedCode.replace("'api.ts'", `'${options.apiDefFilePath}'`));
+                const res = await runTask(wrapCode(apiGlobalName, apiDefFilePath, code));
         
                 if(!res) {
-                    console.error("Task unsuccessful");
+                    console.error("Unknown failure running task");
+                    return {
+                        success: false,
+                        logs: [],
+                    };
                 } 
         
                 while(isAsyncProcessRunning()) {
@@ -219,6 +380,8 @@ async function createWithAPI(options:AIMyAPIOptions): Promise<AIMyAPIInstance> {
                 if(debug) {
                     console.log("Completed");
                 }
+
+                return res;
             } catch(err) {
                 if(debug) {
                     console.error("Error running task:", err);
@@ -228,10 +391,15 @@ async function createWithAPI(options:AIMyAPIOptions): Promise<AIMyAPIInstance> {
             await new Promise((resolve) => setTimeout(resolve, 100));
         
             rejectOpenPromises(vm);
-            vm.dispose()             
-        
+            vm.dispose();
+
+            return {
+                success: false,
+                logs: [],
+            };
         }, 
-        processRequest: async function (userQuery, currentContext:any = undefined):Promise<GenerateCodeResult> {
+        // generates code and runs it in one step
+        processRequest: async function (userQuery: string, currentContext:any = undefined):Promise<GenerateCodeResult | null> {
             // Create prompt
             
             if(debug) {
@@ -240,24 +408,38 @@ async function createWithAPI(options:AIMyAPIOptions): Promise<AIMyAPIInstance> {
         
             // Generate Code
             console.time("Generate Task")
-            const generatedCode:GenerateCodeResult = await generateCode(userQuery, [], createTaskPrompt.replace("{{CONTEXT}}", currentContext ? "```" + JSON.stringify(currentContext, null, 2) + "```" : "" ));
+            const generatedCode:GenerateCodeResult | null = await generateCode(instance, 
+                userQuery, 
+                [], 
+                createTaskPrompt.replace("{{CONTEXT}}", currentContext ? JSON.stringify(currentContext, null, 2) : "" ), 
+                debug, 
+                hideLogsFromAgent, 
+                model, 
+                additionalModelOptions
+            );
             console.timeEnd("Generate Task")
         
             console.time("Run Code")
-            await this.runCode(generatedCode.code.replace("'api.ts'", `'${options.apiDefFilePath}'`));
+            if(generatedCode?.code) {
+                await instance.runCode(generatedCode.code);
+            } else {
+                console.error("No code generated", generatedCode);
+            }
             console.timeEnd("Run Code")
 
             return generatedCode;
         }
     }
 
+    return instance;
+
 }
 
 export interface AIMyAPIModuleExports {
     createWithAPI: (options:AIMyAPIOptions) => Promise<AIMyAPIInstance>;
     
     createBasePrompt: (apiFilePath:string, documentationPath: string) => string;
-    generateCode: (queryText:string, userChatHistory:ChatCompletionMessageParam[], createTaskPrompt:string, debug:boolean) => Promise<GenerateCodeResult>;
+    generateCode: (instance: AIMyAPIInstance, queryText:string, userChatHistory:ChatCompletionMessageParam[], createTaskPrompt:string, debug:boolean) => Promise<GenerateCodeResult | null>;
     createSandbox: (QuickJS:QuickJSWASMModule, globals: any) => Promise<any>;
 }
 
@@ -270,3 +452,6 @@ export const AIMyAPI:AIMyAPIModuleExports = {
     generateCode,
     createSandbox,
 }
+
+export * from "./sandbox";
+