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↗︎