Support code generation with 'define'

.prettierrc.json

@@ -0,0 +1,6 @@
+{
+  "trailingComma": "es5",
+  "tabWidth": 2,
+  "semi": true,
+  "singleQuote": true
+}

README.md

@@ -43,9 +43,7 @@ Add the following snippet to your `tsconfig.json`:
 
 ```json
 "compilerOptions": {
-  "plugins": [
-    { "transform": "ts-askit" }
-  ]
+    "plugins": [{ "transform": "ts-askit/transform" }]
 }
 ```
 

examples/src/pbe.ts

@@ -1,18 +1,19 @@
-import { ask, Example } from 'ts-askit';
+import { define, Example } from 'ts-askit';
 
-function addInBase2(x: string, y: string) {
-  const example: Example[] = [
-    { input: { x: '1', y: '0' }, output: '1' },
-    { input: { x: '1', y: '1' }, output: '10' },
-    { input: { x: '101', y: '11' }, output: '1000' },
-    { input: { x: '1001', y: '110' }, output: '1111' },
-    { input: { x: '1111', y: '1' }, output: '10000' },
-  ];
-  return ask<string>('Add {{x}} and {{y}}', example);
-}
+const example: Example[] = [
+  { input: { x: '1', y: '0' }, output: '1' },
+  { input: { x: '1', y: '1' }, output: '10' },
+  { input: { x: '101', y: '11' }, output: '1000' },
+  { input: { x: '1001', y: '110' }, output: '1111' },
+  { input: { x: '1111', y: '1' }, output: '10000' },
+];
+const addInBase2 = define<string, { x: string; y: string }>(
+  'Add {{x}} and {{y}}',
+  example
+);
 
 async function doit() {
-  console.log(await addInBase2('101', '11'));
+  console.log(await addInBase2({ x: '101', y: '11' }));
 }
 
 doit();

examples/src/top50.ts

@@ -97,7 +97,6 @@ function generateFibonacci(n: number) {
   const example: Example[] = [{ input: { n: 5 }, output: [0, 1, 1, 2, 3] }];
   return ask<number[]>('Generate the Fibonacci sequence up to {{n}}.', example);
 }
-// ... Previous function definitions are correct ...
 
 // Write a function that finds the greatest common divisor of two numbers.
 function findGreatestCommonDivisor(a: number, b: number) {
@@ -163,7 +162,7 @@ function formatDate(d: Date) {
   ];
   return ask<string>(
     'Format the date {{d}} into a specific string format.'
-    //    example
+    /* example */
   );
 }
 
@@ -177,7 +176,7 @@ function findDateDifference(d1: Date, d2: Date) {
   ];
   return ask<number>(
     'Find the difference between the dates {{d1}} and {{d2}}.'
-    //    example
+    /* example */
   );
 }
 
@@ -382,12 +381,3 @@ function findMode(ns: number[]) {
   const example: Example[] = [{ input: { ns: [1, 2, 3, 1, 2, 3] }, output: 1 }];
   return ask<number>('Find the mode of the array {{ns}}.', example);
 }
-
-// Write a function that checks if an array is sorted in ascending order.
-function isSortedAscending(ns: number[]) {
-  const example: Example[] = [{ input: { ns: [1, 2, 3] }, output: true }];
-  return ask<boolean>(
-    'Check if the array {{ns}} is sorted in ascending order.',
-    example
-  );
-}

src/askit.ts

@@ -3,37 +3,60 @@ import { ExamplesType } from './example';
 import { define } from './function';
 import * as t from './types';
 
+let fReason = '';
+let fErrors: string[] = [];
+let fCompletion: any = { data: null };
+
+export function getReason(): string {
+  return fReason;
+}
+export function getErrors(): string[] {
+  return fErrors;
+}
+export function getCompletion(): any {
+  return fCompletion;
+}
+
 export async function llm<T = any>(
   template: string,
-  examples: ExamplesType = []
+  trainingExamples: ExamplesType = [],
+  testExamples: ExamplesType = []
 ): Promise<T> {
-  return ask<T>(template, examples);
+  return ask<T>(template, testExamples, trainingExamples);
 }
 
 export async function ask<T = any>(
   template: string,
-  examples?: ExamplesType
+  trainingExamplesExamples?: ExamplesType,
+  testExamples?: ExamplesType
 ): Promise<T>;
 export async function ask<T>(
   type: t.Type<T>,
   template: string,
-  examples?: ExamplesType,
+  trainingExamplesExamples?: ExamplesType,
+  testExamples?: ExamplesType,
   args?: { [key: string]: any }
 ): Promise<T>;
 export async function ask<T = any>(...args: unknown[]): Promise<T> {
   let type: t.Type<T>;
   let template: string;
-  let examples: ExamplesType;
+  let trainingExamples: ExamplesType;
+  let testExamples: ExamplesType;
 
   if (typeof args[0] === 'string') {
     throw new Error('Transpilation with AskIt is needed!');
   } else if (args[0] instanceof t.Type) {
     type = args[0] as t.Type<T>;
     template = args[1] as string;
-    examples = args[2] as ExamplesType;
+    trainingExamples = args[2] as ExamplesType;
+    testExamples = args[3] as ExamplesType;
     const argsMap = args[3] as { [key: string]: any };
-    const f = define(type, template, examples);
-    const result = f(argsMap);
+    const f = define(type, template, trainingExamples, testExamples);
+    const result = await f(argsMap);
+    fReason = f.reason;
+    fErrors = f.errors;
+    fCompletion = f.completion;
+
     return result;
   }
   throw new Error('Invalid number of arguments');

src/cli/askgen.ts

@@ -232,7 +232,7 @@ if (args.length !== 1) {
   process.exit(1);
 }
 
-function generateFiles(filename: string) {
+async function generateFiles(filename: string) {
   const text = readFileSync(filename, 'utf-8');
   const lines = text.split('\n');
   const info: Info[] = lines.map((line: string) => JSON.parse(line));
@@ -241,7 +241,8 @@ function generateFiles(filename: string) {
   const sourceFilePath = path.join(dirName, sourceFileName);
   // generate source code
   const modules = new Set<string>();
-  info.forEach(async (i: Info) => {
+  for (const i of info) {
+    const start = process.hrtime.bigint();
     const { modulePath } = makeModuleName(sourceFilePath, i.name);
     modules.add(i.name);
     if (fs.existsSync(modulePath)) {
@@ -252,6 +253,7 @@ function generateFiles(filename: string) {
     if (!fs.existsSync(dirPath)) {
       fs.mkdirSync(dirPath, { recursive: true });
     }
+    let testFailedCount = 0;
     for (let count = 0; count < 10; count++) {
       const s = await sendChatRequest(make_message(i.desc, i.params));
       const m = s.match(/\[([0-9])\]/);
@@ -288,19 +290,36 @@ function generateFiles(filename: string) {
           const jsCode = output.outputText;
           // console.log('ts code: ', code);
           // console.log('js code: ', jsCode);
-          let ok = testFunction(jsCode, i);
+          const generatedFunc = getFunction(jsCode, i.name);
+          if (!generatedFunc) {
+            continue;
+          }
+          let ok = await testFunction(generatedFunc, i);
           if (ok) {
-            writeFileSync(modulePath, code);
+            const program = `// Recompilation count: ${testFailedCount}\n${code}`;
+            writeFileSync(modulePath, program);
+            const end = process.hrtime.bigint();
+            const elapsed = Number(end - start) / 1000000000;
             console.log('Generated: ', modulePath, count + 1);
+            writeFileSync(
+              filename + '.log',
+              JSON.stringify({
+                name: i.name,
+                elapsed,
+                count,
+              })
+            );
             break;
+          } else {
+            testFailedCount++;
           }
         } else {
           console.log(diag);
           continue;
         }
       }
     }
-  });
+  }
 
   // delete unused files
   const baseName = path.basename(filename, '.ts.jsonl');
@@ -317,31 +336,60 @@ function generateFiles(filename: string) {
   // });
 }
 
-function testFunction(jsCode: string, i: Info) {
-  let ok = true;
+function getFunction(jsCode: string, functionName: string) {
   try {
     const generatedFunc = eval(jsCode);
-    i.examples.forEach(({ input, output }) => {
-      const args = i.params.map(([_, name]) => input[name]);
-      // const args: string[] = [];
-      // for (let name in inputs) {
-      //   args.push(inputs[name]);
-      // }
-      //console.log('args: ', args);
-      const actualOutput = generatedFunc(...args);
+    if (generatedFunc.name !== functionName) {
+      return false;
+    }
+    return generatedFunc;
+  } catch (e) {
+    return false;
+  }
+}
+
+function executeWithTimeout(
+  func: any,
+  args: any[],
+  timeout: number
+): Promise<any> {
+  return new Promise((resolve, reject) => {
+    const timer = setTimeout(() => {
+      reject(new Error('Timeout'));
+    }, timeout);
+    try {
+      const result = func(...args);
+      clearTimeout(timer);
+      resolve(result);
+    } catch (e) {
+      clearTimeout(timer);
+      reject(e);
+    }
+  });
+}
+
+async function testFunction(generatedFunc: any, i: Info) {
+  let ok = true;
+  try {
+    for (const example of i.testExamples) {
+      const args = i.params.map(([_, name]) => example.input[name]);
+      const actualOutput = await executeWithTimeout(generatedFunc, args, 60000);
       if (
-        !_.isEqual(actualOutput, output) &&
-        !almostEqual(actualOutput, output) // XXX: for float
+        !_.isEqual(actualOutput, example.output) &&
+        !almostEqual(actualOutput, example.output) // XXX: for float
       ) {
-        console.log('Failed: ', input, output, actualOutput, i.desc);
-        console.log('Code: ', jsCode);
+        console.log(
+          'Failed: ',
+          example.input,
+          example.output,
+          actualOutput,
+          i.desc
+        );
         console.log(JSON.stringify(args));
         ok = false;
       }
-    });
+    }
   } catch (e) {
-    console.log(e);
-    console.log(jsCode);
     ok = false;
   }
   return ok;

src/function.ts

@@ -2,7 +2,6 @@ import { ExamplesType, checkExamples } from './example';
 import { generateUniqueFunctionName } from './function-name';
 import { chat } from './gpt';
 import { convertTemplate, extractVariables } from './template';
-import { printType } from './type-printer';
 import * as t from './types';
 
 export class Function<T> {
@@ -37,7 +36,7 @@ export class Function<T> {
     });
   }
 
-  call(args: { [key: string]: any } = {}): Promise<T> {
+  async call(args: { [key: string]: any } = {}) {
     this.checkArgs(args);
     const convertedTemplate = convertTemplate(this.template);
     return chat(
@@ -55,26 +54,48 @@ export class Function<T> {
   }
 }
 
-export function define<T>(
+type DefinedFunctionType<T, U> = {
+  (args: U): Promise<T>;
+  reason: string;
+  errors: string[];
+  completion: any;
+};
+
+export function define<
+  T,
+  U extends Record<string, any> = { [key: string]: any },
+>(
   type: t.Type<T>,
   template: string,
-  trainingExamples?: ExamplesType
-): (args: { [key: string]: any }) => Promise<T>;
-export function define<T>(
+  trainingExamples?: ExamplesType,
+  testExamples?: ExamplesType
+): DefinedFunctionType<T, U>;
+export function define<
+  T,
+  U extends Record<string, any> = { [key: string]: any },
+>(
   template: string,
-  trainingExamples?: ExamplesType
-): (args: { [key: string]: any }) => Promise<T>;
-export function define<T>(
-  ...args: unknown[]
-): (args: { [key: string]: any }) => Promise<T> {
+  trainingExamples?: ExamplesType,
+  testExamples?: ExamplesType
+): DefinedFunctionType<T, U>;
+export function define<
+  T,
+  U extends { [key: string]: any } = { [key: string]: any },
+>(...args: unknown[]): DefinedFunctionType<T, U> {
   if (typeof args[0] != 'string') {
     const type = args[0] as t.Type<T>;
     const template = args[1] as string;
     const trainingExamples = (args[2] as ExamplesType) || [];
     const f = new Function(type, template, trainingExamples);
-    return function (args: { [key: string]: any } = []) {
-      return f.call(args);
-    };
+    const g = async function (args: U = {} as U) {
+      // XXX: as U needed?
+      const [answer, reason, errors, completion] = await f.call(args);
+      g.reason = reason;
+      g.errors = errors;
+      g.completion = completion;
+      return answer;
+    } as DefinedFunctionType<T, U>;
+    return g;
   }
   throw new Error('defined should be transpiled by AskIt');
 }