aimyapi 0.2.1 → 1.0.3

package/dist/index.mjs

@@ -4,6 +4,8 @@ import { OpenAI } from 'openai';
 import path from 'path';
 import { getQuickJS } from 'quickjs-emscripten';
 import * as ts from 'typescript';
+import { z } from 'zod';
+import { zodToJsonSchema } from 'zod-to-json-schema';
 
 var require$1 = (
 			true
@@ -196,8 +198,8 @@ function injectTimingFunctions(vm, cbAddedTimer, cbRemovedTimer, maxTimeout = 6e
   _clearInterval.dispose();
 }
 
-function tsCompile(source, options = null) {
-  if (null === options) {
+function tsCompile(source, options) {
+  if (!options) {
     options = {
       compilerOptions: {
         target: ts.ScriptTarget.ES2020,
@@ -207,6 +209,34 @@ function tsCompile(source, options = null) {
   }
   return ts.transpileModule(source, options).outputText;
 }
+class Logger {
+  constructor(debug = false) {
+    this.logs = [];
+    this.debug = false;
+    this.debug = debug;
+  }
+  log(args, source) {
+    if (Array.isArray(args)) {
+      this.logs.push({ type: "log", args, source });
+    } else {
+      this.logs.push({ type: "log", args: [args], source });
+    }
+    if (this.debug) {
+      console.log(`[${source}]:`, ...args);
+    }
+  }
+  error(args, source) {
+    if (Array.isArray(args)) {
+      this.logs.push({ type: "error", args, source });
+    } else {
+      this.logs.push({ type: "error", args: [args], source });
+    }
+    console.error(`[${source}]:`, ...args);
+  }
+  clear() {
+    this.logs = [];
+  }
+}
 async function createSandbox(QuickJS, requireLookup = {}, globals = {}, debug = false) {
   const vm = QuickJS.newContext();
   let errorState = false;
@@ -223,20 +253,21 @@ async function createSandbox(QuickJS, requireLookup = {}, globals = {}, debug =
   let isAsyncProcessRunning = () => {
     return asyncProcessesRunning > 0;
   };
+  const logger = new Logger(debug);
   const logHandle = vm.newFunction("log", (...args) => {
     const nativeArgs = args.map(vm.dump);
-    console.log("QuickJS Log:", ...nativeArgs);
+    logger.log(nativeArgs, "sandbox");
   });
   const consoleHandle = vm.newObject();
   const exportsHandle = vm.newObject();
   const errorHandle = vm.newFunction("error", (...args) => {
     try {
       const nativeArgs = args.map(vm.dump);
-      console.log("QuickJS console.error:", ...nativeArgs);
+      logger.error(nativeArgs, "sandbox");
       lastError = JSON.stringify(nativeArgs);
       errorState = true;
     } catch (e) {
-      console.log("Error in error:", e);
+      logger.error(["? Error in error:", e], "sandbox");
     }
   });
   const globalNames = Object.getOwnPropertyNames(globals);
@@ -278,170 +309,326 @@ async function createSandbox(QuickJS, requireLookup = {}, globals = {}, debug =
       return lastError;
     },
     isAsyncProcessRunning,
-    runTask: (task) => {
+    // make sure code passes the typescript compilation
+    checkCode: (code) => {
+      logger.clear();
       try {
-        const compiled = tsCompile(task);
+        const compiled = tsCompile(code);
         if (!compiled || compiled.length === 0) {
-          console.error("Could not compile", task);
-          return false;
+          logger.error(["Could not compile:", code], "compiler");
+          return {
+            success: false,
+            logs: logger.logs
+          };
+        }
+      } catch (e) {
+        logger.error(["Error testing code:", e], "compiler");
+        return {
+          success: false,
+          logs: logger.logs
+        };
+      }
+      return {
+        success: true,
+        logs: logger.logs
+      };
+    },
+    runCode: (code) => {
+      try {
+        logger.clear();
+        const compiled = tsCompile(code);
+        if (!compiled || compiled.length === 0) {
+          logger.error(["Could not compile:", code], "compiler");
+          return {
+            success: false,
+            logs: logger.logs
+          };
         }
         errorState = false;
         const result = vm.evalCode(compiled);
         if (result.error) {
-          const lines = compiled.split("\n");
-          for (let i = 0; i < lines.length; i++) {
-            console.log(`${i + 1}: ${lines[i]}`);
-          }
-          console.log("Execution failed:", vm.dump(result.error));
+          const lines = compiled.split("\n").map((line, index) => `${index + 1}: ${line}`);
+          logger.error(["Execution failed:", lines.join("\n"), vm.dump(result.error)], "runner");
           result.error.dispose();
-          return false;
-        } else {
+          return {
+            success: false,
+            logs: logger.logs
+          };
+        } else if ("value" in result) {
           result.value.dispose();
           if (errorState) {
-            const lines = compiled.split("\n");
-            for (let i = 0; i < lines.length; i++) {
-              console.log(`${i + 1}: ${lines[i]}`);
-            }
+            const lines = compiled.split("\n").map((line, index) => `${index + 1}: ${line}`);
+            logger.error(["An error was reported during execution. Compiled code with line numbers:", lines.join("\n")], "runner");
           }
-          return !errorState;
+          return {
+            success: !errorState,
+            logs: logger.logs
+          };
+        } else {
+          logger.error(["Unknown execution result:", result], "runner");
+          return {
+            success: false,
+            logs: logger.logs
+          };
         }
       } catch (e) {
-        console.log(e);
+        logger.error(["Error running code:", e], "runner");
       }
-      return false;
+      return {
+        success: false,
+        logs: logger.logs
+      };
     }
   };
 }
 
 require$1("dotenv").config();
-const openai = new OpenAI({
-  apiKey: process.env.OPENAI_API_KEY
-  // This is also the default, can be omitted
-});
-const CODE_START = `\`\`\`
-import * as ApiDefs from 'api.ts'
-
-(async() {
-    try {`;
-const CODE_END = `    } catch(err) { 
+const openai = new OpenAI();
+let QuickJS = null;
+const wrapCode = (apiGlobalName, apiFileDefPath, code) => `
+(async () => {
+    try {
+        ${code}
+    } catch(err) { 
         console.error(err); 
     }
 })();        
-\`\`\``;
-function indentString(str, numTabs) {
-  const tabs = "	".repeat(numTabs);
-  return str.split("\n").map((line) => tabs + line).join("\n");
+`;
+const wrapCodeBlock = (code, codeType = "") => `
+\`\`\`${codeType}
+${code}
+\`\`\`
+`;
+function zodParseJSON(schema) {
+  return (input) => schema.parse(JSON.parse(input));
 }
-function replaceCodeBlocks(text) {
-  const codeBlockRegex = /```[\s\S]*?```/g;
-  return text.replace(codeBlockRegex, (match) => {
-    const code = match.slice(3, -3).trim();
-    return `${CODE_START}
-${indentString(code, 2)}
-${CODE_END}`;
-  });
-}
-function createBasePrompt(apiFilePath, documentationPath) {
+function createBasePrompt(apiFilePath, apiGlobalName, documentationPath) {
   const apiText = fs.readFileSync(apiFilePath, "utf8");
   const documentationText = documentationPath ? fs.readFileSync(documentationPath, "utf8") : "";
-  const createTaskPrompt = fs.readFileSync(path.join(__dirname, "../prompts/create-task-prompt.md"), "utf8").replace("{{DOCUMENTATION}}", replaceCodeBlocks(documentationText)).replace("{{API}}", apiText);
+  const createTaskPrompt = fs.readFileSync(path.join(__dirname, "../prompts/create-task-prompt.md"), "utf8").replace("{{DOCUMENTATION}}", documentationText).replace("{{API}}", apiText).replace("{{API_GLOBAL_NAME}}", apiGlobalName || "");
   return createTaskPrompt;
 }
-const generateCode = async function(queryText, userChatHistory, createTaskPrompt, debug = false, model = "gpt-4o-mini") {
-  if (!queryText)
-    return { code: "", loggableCode: "" };
-  const prompt = createTaskPrompt.replace("{{QUERY_TEXT}}", queryText);
-  const messages = [
+const runCodeSchema = z.object({
+  code: z.string()
+});
+const checkCodeSchema = z.object({
+  code: z.string()
+});
+const submitCodeSchema = z.object({
+  code: z.string()
+});
+const generateCode = async function(instance, queryText, userChatHistory, createTaskPrompt, debug = false, hideLogsFromAgent = false, model = "gpt-5-mini", additionalModelOptions) {
+  let finalSubmittedCode = "";
+  if (debug) {
+    console.log("Generating code for query:", queryText);
+  }
+  const codingTools = [
     {
-      role: "system",
-      content: prompt
+      type: "function",
+      function: {
+        name: "runCode",
+        // @ts-ignore
+        function: async ({ code }) => {
+          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)
+      }
+    },
+    {
+      type: "function",
+      function: {
+        name: "checkCode",
+        // @ts-ignore
+        function: async ({ code }) => {
+          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)
+      }
     },
-    ...userChatHistory,
     {
-      role: "user",
-      content: `${queryText}. Respond with json`,
-      name: "user"
+      type: "function",
+      function: {
+        name: "submitCode",
+        // @ts-ignore
+        function: async ({ code }) => {
+          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,
-        temperature: 0,
-        max_tokens: 2e3,
-        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) => {
+      if (debug) {
+        console.log("Message:", JSON.stringify(message, null, 2));
       }
-    );
-    if (debug)
-      console.log(results.usage);
-    const response = results.choices[0].message.content;
-    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: "" };
+    }).on("functionToolCall", (functionCall) => {
+      if (debug) {
+        console.log("Function Call:", JSON.stringify(functionCall, null, 2));
+      }
+    }).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 || void 0
+    };
   } catch (err) {
     console.error(err);
-    return { code: "", loggableCode: "" };
+    return null;
   }
 };
-let QuickJS = null;
 async function createWithAPI(options) {
-  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
   };
-  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 createTaskPrompt = createBasePrompt(apiDefFilePath, apiDocsPath);
-  return {
-    options,
+  const { apiObject, apiWhitelist, apiGlobalName, apiExports, apiDefFilePath, apiDocsPath, debug, model, additionalModelOptions, hideLogsFromAgent } = localOptions;
+  const createTaskPrompt = createBasePrompt(apiDefFilePath, apiGlobalName, apiDocsPath);
+  const instance = {
+    options: localOptions,
     generateCode: async function(queryText, userChatHistory, currentContext = void 0) {
       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) {
+    checkCode: async function(generatedCode) {
       if (!QuickJS) {
         if (debug)
           console.log("Creating quick.js");
         QuickJS = await getQuickJS();
       }
       let apiGlobals = {};
-      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[key]
+          };
+        }
+      }
+      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: ["\u2714\uFE0F Code check passed"], source: "compiler" }, ...res.logs]
+          };
+        }
+        return {
+          success: false,
+          logs: [{ type: "error", args: ["\u274C 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) {
+      if (!QuickJS) {
+        if (debug)
+          console.log("Creating quick.js");
+        QuickJS = await getQuickJS();
       }
-      const { vm, runTask, isAsyncProcessRunning } = await createSandbox(QuickJS, {
-        [apiDefFilePath.replaceAll("\\", "")]: apiExports
+      let apiGlobals = {};
+      if (options.apiGlobals) {
+        for (const key in options.apiGlobals) {
+          apiGlobals[key] = {
+            value: options.apiGlobals[key]
+          };
+        }
+      }
+      const { vm, runCode: runTask, isAsyncProcessRunning } = await createSandbox(QuickJS, {
+        [apiDefFilePath]: apiExports
       }, {
         ...apiGlobals,
         [apiGlobalName]: {
@@ -451,11 +638,15 @@ async function createWithAPI(options) {
       }, options.debug);
       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()) {
           await new Promise((resolve) => setTimeout(resolve, 100));
@@ -463,6 +654,7 @@ async function createWithAPI(options) {
         if (debug) {
           console.log("Completed");
         }
+        return res;
       } catch (err) {
         if (debug) {
           console.error("Error running task:", err);
@@ -471,20 +663,39 @@ async function createWithAPI(options) {
       await new Promise((resolve) => setTimeout(resolve, 100));
       rejectOpenPromises(vm);
       vm.dispose();
+      return {
+        success: false,
+        logs: []
+      };
     },
+    // generates code and runs it in one step
     processRequest: async function(userQuery, currentContext = void 0) {
       if (debug) {
         console.log("Query:", userQuery);
       }
       console.time("Generate Task");
-      const generatedCode = await generateCode(userQuery, [], createTaskPrompt.replace("{{CONTEXT}}", currentContext ? "```" + JSON.stringify(currentContext, null, 2) + "```" : ""));
+      const generatedCode = 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;
 }
 const AIMyAPI = {
   // Standard way of using the library
@@ -495,4 +706,4 @@ const AIMyAPI = {
   createSandbox
 };
 
-export { AIMyAPI, createBasePrompt, generateCode };
+export { AIMyAPI, createBasePrompt, generateCode, zodParseJSON };

package/example/example1.ts

@@ -1,6 +1,5 @@
 import path from "path";
 import { AIMyAPI } from "../src";
-import * as APIExports from "./sales_api";
 import { API } from "./sales_api_impl";
 
 // Simpler example using no support for chat history; 
@@ -10,21 +9,23 @@ import { API } from "./sales_api_impl";
 
     const aimyapi = await AIMyAPI.createWithAPI({
         apiObject: api,
-        apiGlobalName: "api", // Should match whatever you declared as your global in your api.
-        apiExports: APIExports,
+        apiGlobalName: "salesApi", // Should match whatever you declared as your global in your api.
         apiDefFilePath: path.join(__dirname, "./sales_api.ts"),
         apiDocsPath: path.join(__dirname, "./sales_api.md"),
-        debug: true,
+        debug: false,
     })
+
+    aimyapi.runCode(`console.log("Hello from the sandbox! Here is some sales data:", salesApi.getSalesData())`);
     
-    // console.log("Print 'hello' three times.")
-    // await aimyapi.processRequest("Print 'hello' three times.");
+    console.log("Print 'hello' three times.")
+    await aimyapi.processRequest("Print 'hello' three times.");
 
-    // console.log("What month had the highest sales in 2020 for any business unit?")
-    // await aimyapi.processRequest("What was our best month in 2020");
+    console.log("What month had the highest sales in 2020 for any business unit?")
+    await aimyapi.processRequest("What was our best month in 2020");
 
-    // console.log("Which business unit had the highest sales in 2020?")
-    // await aimyapi.processRequest("Which business unit had the highest sales in 2020?");
+    console.log("Which business unit had the highest sales in 2020? Print the totals for each")
+    await aimyapi.processRequest("Which business unit had the highest sales in 2020? Print the totals for each");
+    // Note: this is a bit of a trick question because two of them are tied for first. The agent usually doesn't realize this edge case. Should we instruct it to look for edge cases?
 
     console.log("Compute the total sales for each month in 2020 then email the top 3 to test@email.com")
     await aimyapi.processRequest("Compute the total sales for each month in 2020 then email the top 3 to test@email.com");