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

[rajatsandeepsen]

Proposal for generateCode API

Imagine the LLM can write Javascript code with custom logic within the help of limited tools inside a safe eval( )

Usage

Let me show you an example on how to build an AI powered banking app

import { z } from "zod"
import { tool, experimental_generateCode as generateCode } from "ai"
import { model } from "./model"

Define a set of tools

let balance = 30 // simulation of a DB
const history = [
    { amount: 10, from: "Bob" },
    { amount: 20, to: "Alice" },
]

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 })
        },
        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().optional(), from: z.string().optional() }))
    })
})

Fun part begins here

const result = await generateCode({
    model,
    system: "You are a banking app",
    tools: tools,
    prompt: "Get history and find amount i got from Bob, then send that amount to Bob. Then again get history and balance",
})

console.log("Code:", result.code)
console.log("Schema:", result.schema)
console.log("Output:", await result.execute())

This is the output after execution await result.execute( )

{
  balance: 20,
  history: [
    {
      amount: 20,
      to: "Alice"
    }, {
      amount: 10,
      from: "Bob"
    },{
      amount: 10,
      to "Bob"
    }
  ]
}

This is the code written by LLM result.code

let history = this.getHistory({});
let amountGotFromBob = 0;
for (let transaction of history) {
    if (transaction?.from === 'Bob') {
        amountGotFromBob += transaction.amount;
    }
}
if (amountGotFromBob > 0) {
    this.sentMoney({ amount: amountGotFromBob, receiver: 'Bob' });
}
let balance = this.getBalance({});
return { balance, history };

This is the JSON schema written by LLM result.schema (useful for generative UI)

{
  $schema: "http://json-schema.org/draft-07/schema#",
  type: "object",
  properties: {
    balance: {
      type: "number",
    },
    history: {
      type: "array",
      items: {
        type: "object",
        properties: {
          amount: {
            type: "number",
          },
          to: {
            type: "string",
          },
          from: {
            type: "string",
          },
        },
        required: [ "amount" ],
      },
    },
  },
  required: [ "balance", "history" ]
}

Instead of multi-step toolResults, now LLM can write logic along with tools provided by developer. Also the LLM can't execute malicious code because of a safety simple technique i implemented in this.

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, [])
}
const main = (code:string) => `const main = async () => {\n${code}\n}\nreturn main()`

const generateCode = ({ ... }) => {

	// some generateText logic

	return {
		code, schema,
		execute: createFunction(tools, code)
	}
}

Therefor this LLM is restricted to only invoke list of functions we provided to them.

The generateCode( ) API is a powerful wrapper around generateText( )

[socket-security]

New dependencies detected. Learn more about Socket for GitHub ↗︎

Package	New capabilities	Transitives	Size	Publisher
npm/[email protected]	None	0	34.6 kB	sachinraja

View full report↗︎

[lgrammel]

This is pretty cool. However, there is one big issue: it includes a fairly large prompt. The AI SDK tries to not include any prompts whenever possible, because they end up being provider and model dependent (there is some minimal prompting in JSON generation but that is all). Do you see ways to do this without prompting? Or could there be alternative approaches that minimize it?

[lgrammel]

Removing the optionality here will break important functionality, namely tools without execute.

[rajatsandeepsen]

I got three alternatives idea for this issue

• Create an isolated tool in different name (eg: etool specifically for generateCode

• add one more option to toolChoice: "none" | "auto" | "required" | "code" (this code helps to make execute strictly required

• throw error is execute param is undefined

Give me a feedback, I'll rework on this

[lgrammel]

RESULT extends Parameters seems strange. Is this intentional?

[rajatsandeepsen]

generateCode requires an returns zodSchema so LLM can understand whats the output of each tools.

It'll help LLM to create variables to save the tool-calling-result and pass it to next tool.

Reason why i used Parameters type is to strictly specify execute tool according to its returns: zodSchema

I could rework this as

type Returns = Parameters // same properties

export function tool<PARAMETERS extends Parameters, RESULT, RETURNS extends Returns>(
tool: CoreTool<PARAMETERS, RETURNS> & {
    execute: (
      args: inferParameters<PARAMETERS>,
      options: ToolExecutionOptions,
    ) => PromiseLike<inferParameters<RETURNS>>;
  }
)

[rajatsandeepsen]

This is pretty cool. However, there is one big issue: it includes a fairly large prompt. The AI SDK tries to not include any prompts whenever possible, because they end up being provider and model dependent (there is some minimal prompting in JSON generation but that is all). Do you see ways to do this without prompting? Or could there be alternative approaches that minimize it?

This idea of tools inside generateCode is not widely adopted in openAI API specs.

Everyone uses JSON schema for specifying tools to LLM (as Industry standards). But i used zod-to-ts for three reasons

1. it converts zodSchema to typescript definition
2. takes less tokens for defining same functionality compared to JSON schema
3. anyway we're not passing these functions to tools part of OpenAI API and not returning a list of tool name in JSON object

So complex parameters & returns type can be specified + it can document each param usage like

const writeLetter = tool({
	description: "write letter to a person",
	parameters: z.object({
	 name: z.string().describe("Name of the person")
	}),
	returns: z.string(),
	execute: ({ name }) => `Hello ${name}`
})

After zod-to-ts

const writeLetter = (params: {
	// Name of the person
	name: stirng
}) :string => {
	// some logic
}

Prompt i using is to ensure LLM can writes better and secure code. Some of these prompts are important.