Skip to content

Commit

Permalink
fix: pnpm-lock file & formatting
Browse files Browse the repository at this point in the history
  • Loading branch information
rajatsandeepsen committed Dec 27, 2024
1 parent 7fbfcdd commit dd58085
Show file tree
Hide file tree
Showing 8 changed files with 240 additions and 202 deletions.
19 changes: 13 additions & 6 deletions packages/ai/core/generate-code/function.ts
Original file line number Diff line number Diff line change
@@ -1,12 +1,19 @@
import type { CoreTool } from "../tool"
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 }), {})
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, [])
}
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()`
const main = (code: string) =>
`const main = async () => {\n${code}\n}\nreturn main()`;
88 changes: 43 additions & 45 deletions packages/ai/core/generate-code/generate-code.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,53 +10,52 @@ const dummyResponseValues = {
usage: { promptTokens: 10, completionTokens: 20 },
};

let balance = 30
let balance = 30;
const history = [
{ amount: 10, from: "Bob", to: "me" },
{ amount: 20, to: "Alice", from: "me" },
]
{ 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
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() }))
})
})
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: `\`\`\`js
Expand All @@ -71,18 +70,17 @@ return balance
\`\`\`
`,
};

},
}),
prompt: "Get my balance",
system: "You are a banking app",
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.code, `let balance = this.getBalance({});
// return balance`);
// assert.deepStrictEqual(result.schema, `{
// "type": "number"
// }`);
assert.deepStrictEqual(result.execute(), 30);
});
});
120 changes: 63 additions & 57 deletions packages/ai/core/generate-code/generate-code.ts
Original file line number Diff line number Diff line change
@@ -1,15 +1,20 @@
import { generateText } from "../generate-text/index"
import { newSystemPrompt } from "./prompt"
import { createFunction } from "./function"
import type { CoreTool } from "../tool"
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 js_regex = /```js\n([\s\S]*?)\n```/g;
const json_regex = /```json\n([\s\S]*?)\n```/g;

const thisKeyWord = "listOfFunctions"
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 }
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.
Expand All @@ -19,54 +24,55 @@ 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 {
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: undefined,
schema,
code,
execute: evalCode,
}
}

export { generateCode }
};
} catch (e) {}
}

return {
...result,
schema: undefined,
code,
execute: evalCode,
};
};

export { generateCode };
2 changes: 1 addition & 1 deletion packages/ai/core/generate-code/index.ts
Original file line number Diff line number Diff line change
@@ -1 +1 @@
export { generateCode as experimental_generateCode } from './generate-code';
export { generateCode as experimental_generateCode } from './generate-code';
47 changes: 31 additions & 16 deletions packages/ai/core/generate-code/prompt.ts
Original file line number Diff line number Diff line change
@@ -1,36 +1,51 @@
import type { CoreTool } from "../tool";
import { generateZodTypeString } from "./type-zod";
import type { CoreTool } from '../tool';
import { generateZodTypeString } from './type-zod';

type AnyFunction = (...args: any[]) => any;

function checkAsync(fn: AnyFunction): boolean {
return fn.constructor.name === 'AsyncFunction';
}

type TOOLS = Record<string, CoreTool>
type TOOLS = Record<string, CoreTool>;

const functionDefinition = (tool: TOOLS[string], isAsync:boolean) => {
const type = generateZodTypeString(tool.returns, "returns")
const functionDefinition = (tool: TOOLS[string], isAsync: boolean) => {
const type = generateZodTypeString(tool.returns, 'returns');

const returnType = isAsync ? `Promise<${type}>` : type
const paramType = generateZodTypeString(tool.parameters, "data")
const returnType = isAsync ? `Promise<${type}>` : type;
const paramType = generateZodTypeString(tool.parameters, 'data');

return `${isAsync ? "async" : ""} (data:${paramType}): ${returnType} => {\n\t// ${tool?.description ?? "Does something"}\n\treturn // something\n}`
}
return `${
isAsync ? 'async' : ''
} (data:${paramType}): ${returnType} => {\n\t// ${
tool?.description ?? 'Does something'
}\n\treturn // something\n}`;
};

const displayToolsToCode = (tools: TOOLS) =>
Object.entries(tools)
.map(([toolName, toolObject]) => {
if (!toolObject.execute)
throw new Error(`Execute function is required for tool ${toolName}`);
throw new Error(
`Execute function of tool "${toolName}" is not specified`,
);

if (!toolObject.returns)
throw new Error(
`Return Zod schema of tool "${toolName}" is not specified`,
);

const isAsync = checkAsync(toolObject.execute)
const isAsync = checkAsync(toolObject.execute);

return `const ${toolName} = ${functionDefinition(toolObject, isAsync)}`
return `const ${toolName} = ${functionDefinition(toolObject, isAsync)}`;
})
.join("\n\n")
.join('\n\n');

export const newSystemPrompt = (text: string, tools: TOOLS, thisKeyWord: string) => `Your Persona: ${text}
export const newSystemPrompt = (
text: string,
tools: TOOLS,
thisKeyWord: string,
) => `Your Persona: ${text}
Instructions:
- write pure javascript code
Expand All @@ -52,8 +67,8 @@ Tools:
${displayToolsToCode(tools)}
const ${thisKeyWord} = {
${Object.keys(tools).join(", ")}
${Object.keys(tools).join(', ')}
}
Using above functions, write code to solve the user prompt
`
`;
8 changes: 4 additions & 4 deletions packages/ai/core/generate-code/type-zod.ts
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import type { ZodSchema } from "zod"
import { printNode, zodToTs } from "zod-to-ts"
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
};
return typeString;
};
Loading

0 comments on commit dd58085

Please sign in to comment.