feat (ai/core): Proposal for generate code API with tools

packages/ai/core/generate-code/function.ts

@@ -0,0 +1,12 @@
+import type { CoreTool } from "../tool"
+
+// Creates a safe execution ground for the code
+export const createFunction = (tools: Record<string, CoreTool>, code: string) => {
+  const data = Object.entries(tools).reduce((acc, [key, value]) => ({ ...acc, [key]: value.execute }), {})
+
+  return async () => await new Function(main(code)).apply(data, [])
+}
+
+// Don't remove this.
+// This is the only reason why async and sync function works inside `new Function()`
+const main = (code:string) => `const main = async () => {\n${code}\n}\nreturn main()`

packages/ai/core/generate-code/generate-code.test.ts

@@ -0,0 +1,88 @@
+import assert from 'node:assert';
+import { z } from 'zod';
+import { MockLanguageModelV1 } from '../test/mock-language-model-v1';
+import { generateCode } from './generate-code';
+import { tool } from '../tool';
+
+const dummyResponseValues = {
+  rawCall: { rawPrompt: 'prompt', rawSettings: {} },
+  finishReason: 'stop' as const,
+  usage: { promptTokens: 10, completionTokens: 20 },
+};
+
+let balance = 30
+const history = [
+    { amount: 10, from: "Bob", to: "me" },
+    { amount: 20, to: "Alice", from: "me" },
+]
+
+const tools = ({
+    getBalance: tool({
+        description: "get balance of the user",
+        parameters: z.object({}),
+        execute: async () => {
+            return balance
+        },
+        returns: z.number()
+    }),
+    sentMoney: tool({
+        description: "send money to the user",
+        parameters: z.object({ amount: z.number(), receiver: z.string() }),
+        execute: async ({ amount, receiver }) => {
+            if (balance < amount) {
+                throw new Error("Insufficient balance")
+            }
+            balance -= amount
+
+            history.push({ amount, to: receiver, from: "me" })
+        },
+        returns: z.void()
+    }),
+    getHistory: tool({
+        description: "get history of transactions",
+        parameters: z.unknown(),
+        execute: async () => {
+            return history
+        },
+        returns: z.array(
+            z.object({ amount: z.number(), to: z.string(), from: z.string() }))
+    })
+})
+
+describe('result.code', () => {
+  it('should generate code', async () => {
+
+    const result = await generateCode({
+      tools,
+      model: new MockLanguageModelV1({
+        doGenerate: async ({ prompt, mode }) => {
+
+          return {
+            ...dummyResponseValues,
+            text: `\`\`\`ts
+let balance = listOfFunctions.getBalance({});
+return balance
+\`\`\`
+
+\`\`\`json
+{
+  "type": "number"
+}
+\`\`\`
+`,
+          };
+
+        },
+      }),
+      prompt: "Get my balance",
+      system: "You are a banking app",
+    });
+
+//     assert.deepStrictEqual(result.code, `let balance = this.getBalance({});
+// return balance`);
+//     assert.deepStrictEqual(result.schema, `{
+//   "type": "number"
+// }`);
+    assert.deepStrictEqual(result.execute(), 30);
+  });
+});

packages/ai/core/generate-code/generate-code.ts

@@ -0,0 +1,72 @@
+import { generateText } from "../generate-text/index"
+import { newSystemPrompt } from "./prompt"
+import { createFunction } from "./function"
+import type { CoreTool } from "../tool"
+
+const js_regex = /```js\n([\s\S]*?)\n```/g
+const json_regex = /```json\n([\s\S]*?)\n```/g
+
+const thisKeyWord = "listOfFunctions"
+
+type GenerateCodeParams = Omit<Parameters<typeof generateText>[0], "tools"> & { tools: Record<string, CoreTool> }
+type GenerateCodeReturns = Omit<Awaited<ReturnType<typeof generateText>>, "toolCalls" | "toolResults" | "steps"> & { code: string, execute: (() => Promise<unknown>), schema?: object }
+
+/**
+Generate code that can be executed, un-typed result but with JSON schema for a given prompt and tools using a language model.
+
+This function does not stream the output.
+
+@returns
+A result object that contains the generated code, JSON schema, executable function, the finish reason, the token usage, and additional information.
+ */
+const generateCode = async ({ tools, system, ...rest }: GenerateCodeParams): Promise<GenerateCodeReturns> => {
+
+    const systemNew = newSystemPrompt(
+        system ?? "Follow the instructions and write code for the prompt",
+        tools,
+        thisKeyWord
+    )
+
+    const result = await generateText({
+        ...rest,
+        toolChoice: "none",
+        system: systemNew
+    })
+
+    const codeBlock = result.text.match(js_regex)
+    const jsonCodeBlock = result.text.match(json_regex)
+
+    if (!codeBlock) {
+        throw new Error("No code block found")
+    }
+
+
+    const code = codeBlock[0].replace(/```js\n|```/g, "")
+        .replaceAll(thisKeyWord, "this")
+
+    const evalCode = createFunction(tools, code)
+
+    if (jsonCodeBlock) {
+        try {
+            const schema = JSON.parse(jsonCodeBlock[0]
+                .replace(/```json\n|```/g, ""))
+
+            return {
+                ...result,
+                schema,
+                code,
+                execute: evalCode,
+            }
+        }
+        catch (e) { }
+    }
+
+    return {
+        ...result,
+        schema: undefined,
+        code,
+        execute: evalCode,
+    }
+}
+
+export { generateCode }

packages/ai/core/generate-code/index.ts

@@ -0,0 +1 @@
+export { generateCode as experimental_generateCode } from './generate-code';

packages/ai/core/generate-code/prompt.ts

@@ -0,0 +1,52 @@
+import type { CoreTool } from "../tool";
+import { generateZodTypeString } from "./type-zod";
+
+type AnyFunction = (...args: any[]) => any;
+
+function isAsync(fn: AnyFunction): boolean {
+  return fn.constructor.name === 'AsyncFunction';
+}
+
+type TOOLS = Record<string, CoreTool>
+
+const displayToolsToType = (tools: TOOLS) =>
+  Object.entries(tools)
+    .map(([key, value]) => {
+      const async = isAsync(value.execute)
+      return `type ${key} = (data:${generateZodTypeString(value.parameters, "data")}) => ${async ? "Promise<" : ""}${generateZodTypeString(value.returns, "returns")}${async ? ">" : ""}`
+    }
+    ).join("\n")
+
+const displayToolsToCode = (tools: TOOLS) =>
+  Object.entries(tools)
+    .map(([key, value]) => `const ${key} = ${isAsync(value.execute) ? "async" : ""}(data:${generateZodTypeString(value.parameters, "data")}):${generateZodTypeString(value.returns, "returns")} => {\n    // ${value?.description ?? "Does something"}\n    return // something\n}`)
+    .join("\n\n")
+
+export const newSystemPrompt = (text: string, tools: TOOLS, thisKeyWord: string) => `Your Persona: ${text}
+
+Instructions:
+- write pure javascript code
+- only use functions from the "Tools" list
+- functions are already defined
+- don't imported or redifined
+- nested functions are allowed
+- don't use any external libraries
+- don't use console.log
+- don't wrap code in a function
+- use let to declare variables
+- always end the code with return statement
+- wrap the entire Javascript code in \`\`\`js ... \`\`\` code block
+- also write the JSON schema for the return value of the code in \`\`\`json ... \`\`\` code block
+
+if function name is build(), then use it as ${thisKeyWord}.build()
+
+
+Tools:
+${displayToolsToCode(tools)}
+
+const ${thisKeyWord} = {
+    ${Object.keys(tools).join(", ")}
+}
+
+Using above functions, write code to solve the user prompt
+`

packages/ai/core/generate-code/type-zod.ts

@@ -0,0 +1,9 @@
+import type { ZodSchema } from "zod"
+import { printNode, zodToTs } from "zod-to-ts"
+
+// This function is used to convert zod schema to type definition string
+export const generateZodTypeString = (zod: ZodSchema, K: string) => {
+  const { node: type } = zodToTs(zod, K as string);
+  const typeString = printNode(type);
+  return typeString
+};

packages/ai/core/index.ts

@@ -4,6 +4,7 @@ export * from './data-stream';
 export * from './embed';
 export * from './generate-image';
 export * from './generate-object';
+export * from './generate-code';
 export * from './generate-text';
 export * from './middleware';
 export * from './prompt';

packages/ai/core/tool/tool.ts

@@ -36,14 +36,21 @@ This enables the language model to generate the input.
 
 The tool can also contain an optional execute function for the actual execution function of the tool.
  */
-export type CoreTool<PARAMETERS extends Parameters = any, RESULT = any> = {
+export type CoreTool<PARAMETERS extends Parameters = any, RESULT extends Parameters = any> = {
   /**
 The schema of the input that the tool expects. The language model will use this to generate the input.
 It is also used to validate the output of the language model.
 Use descriptions to make the input understandable for the language model.
    */
   parameters: PARAMETERS;
 
+  returns: RESULT;
+
+  /**
+An optional description of what the tool does. Will be used by the language model to decide whether to use the tool.
+   */
+description?: string;
+
   /**
 Optional conversion function that maps the tool result to multi-part tool content for LLMs.
    */
@@ -56,21 +63,16 @@ If not provided, the tool will not be executed automatically.
 @args is the input of the tool call.
 @options.abortSignal is a signal that can be used to abort the tool call.
    */
-  execute?: (
+  execute: (
     args: inferParameters<PARAMETERS>,
     options: ToolExecutionOptions,
-  ) => PromiseLike<RESULT>;
+  ) => PromiseLike<inferParameters<RESULT>>;
 } & (
   | {
       /**
 Function tool.
        */
       type?: undefined | 'function';
-
-      /**
-An optional description of what the tool does. Will be used by the language model to decide whether to use the tool.
-   */
-      description?: string;
     }
   | {
       /**
@@ -94,28 +96,28 @@ The arguments for configuring the tool. Must match the expected arguments define
 Helper function for inferring the execute args of a tool.
  */
 // Note: special type inference is needed for the execute function args to make sure they are inferred correctly.
-export function tool<PARAMETERS extends Parameters, RESULT>(
+export function tool<PARAMETERS extends Parameters, RESULT extends Parameters>(
   tool: CoreTool<PARAMETERS, RESULT> & {
     execute: (
       args: inferParameters<PARAMETERS>,
       options: ToolExecutionOptions,
-    ) => PromiseLike<RESULT>;
+    ) => PromiseLike<inferParameters<RESULT>>;
   },
 ): CoreTool<PARAMETERS, RESULT> & {
   execute: (
     args: inferParameters<PARAMETERS>,
     options: ToolExecutionOptions,
-  ) => PromiseLike<RESULT>;
+  ) => PromiseLike<inferParameters<RESULT>>;
 };
-export function tool<PARAMETERS extends Parameters, RESULT>(
+export function tool<PARAMETERS extends Parameters, RESULT extends Parameters>(
   tool: CoreTool<PARAMETERS, RESULT> & {
     execute?: undefined;
   },
 ): CoreTool<PARAMETERS, RESULT> & {
   execute: undefined;
 };
-export function tool<PARAMETERS extends Parameters, RESULT = any>(
+export function tool<PARAMETERS extends Parameters, RESULT extends Parameters = any>(
   tool: CoreTool<PARAMETERS, RESULT>,
 ): CoreTool<PARAMETERS, RESULT> {
   return tool;
-}
+}

packages/ai/package.json

@@ -65,7 +65,8 @@
     "@ai-sdk/ui-utils": "1.0.6",
     "@opentelemetry/api": "1.9.0",
     "jsondiffpatch": "0.6.0",
-    "zod-to-json-schema": "^3.23.5"
+    "zod-to-json-schema": "^3.23.5",
+    "zod-to-ts": "^1.2.0"
   },
   "devDependencies": {
     "@edge-runtime/vm": "^5.0.0",