gristlabs_grist-core/test/formula-dataset/runCompletion_impl.ts
2024-04-29 14:54:36 -04:00

360 lines
12 KiB
TypeScript

/**
* This module holds an evaluation scripts for AI assistance. It tests ai assistance on the formula
* dataset. The formula dataset is made of an index file (formula-dataset-index.csv) and a list of
* grist documents hosted on S3. A row in the index file, reference one column (doc_id, table_id,
* col_id) amongst theses documents and a free-text description.
*
* For each entries of the data set, the scripts load the document, requests assistance based on the
* description, and applies the suggested actions to the document. Then it compares the col values
* before and after. Finally it reverts the modification.
*
* The list of grist documents for the formula dataset is a screenshot of all templates document
* taken somewhere in the beginning of Feb 2023.
*
* The script maintains a simple cache of all request to AI to save on the ai requests.
*
* USAGE:
* OPENAI_API_KEY=<my_openai_api_key> node core/test/formula-dataset/runCompletion.js
* or
* ASSISTANT_CHAT_COMPLETION_ENDPOINT=http.... node core/test/formula-dataset/runCompletion.js
* (see Assistance.ts for more options).
*
* # WITH VERBOSE:
* VERBOSE=1 OPENAI_API_KEY=<my_openai_api_key> node core/test/formula-dataset/runCompletion.js
*
* # to reset cache
* rm core/test/formula-dataset/data/cache.json
*/
import { ActiveDoc, Deps as ActiveDocDeps } from "app/server/lib/ActiveDoc";
import { DEPS, sendForCompletion } from "app/server/lib/Assistance";
import log from 'app/server/lib/log';
import crypto from 'crypto';
import { parse } from 'csv-parse/sync';
import fetch, {RequestInfo, RequestInit, Response} from 'node-fetch';
import * as fs from "fs";
import JSZip from "jszip";
import { isEqual, MapCache } from "lodash";
import path from 'path';
import * as os from 'os';
import { pipeline } from 'stream';
import { createDocTools } from "test/server/docTools";
import { promisify } from 'util';
import { AssistanceResponse, AssistanceState } from "app/common/AssistancePrompts";
import { CellValue } from "app/plugin/GristData";
const streamPipeline = promisify(pipeline);
const DATA_PATH = process.env.DATA_PATH || path.join(__dirname, 'data');
const PATH_TO_DOC = path.join(DATA_PATH, 'templates');
const PATH_TO_RESULTS = path.join(DATA_PATH, 'results');
const PATH_TO_CSV = path.join(DATA_PATH, 'formula-dataset-index.csv');
const PATH_TO_CACHE = path.join(DATA_PATH, 'cache');
const TEMPLATE_URL = "https://grist-static.com/datasets/grist_dataset_formulai_2023_02_20.zip";
const oldFetch = DEPS.fetch;
interface FormulaRec {
no_formula: string;
table_id: string;
col_id: string;
doc_id: string;
Description: string;
}
const _stats = {
callCount: 0,
};
const SIMULATE_CONVERSATION = true;
const FOLLOWUP_EVALUATE = false;
export async function runCompletion() {
// This could take a long time for LLMs running on underpowered hardware >:)
ActiveDocDeps.ACTIVEDOC_TIMEOUT = 500000;
// if template directory not exists, make it
if (!fs.existsSync(path.join(PATH_TO_DOC))) {
fs.mkdirSync(path.join(PATH_TO_DOC), {recursive: true});
// create tempdir
const dir = fs.mkdtempSync(path.join(os.tmpdir(), 'grist-templates-'));
const destPath = path.join(dir, 'template.zip');
// start downloading
console.log(
`source url: ${TEMPLATE_URL}\n` +
`destination: ${destPath}\n` +
`download...`
);
const response = await fetch(TEMPLATE_URL);
if (!response.ok) { throw new Error(`unexpected response ${response.statusText}`); }
await streamPipeline(response.body, fs.createWriteStream(destPath));
console.log('done!\n\n' +
'start extraction...');
// unzip to directory
const data = fs.readFileSync(destPath);
const zip = await JSZip.loadAsync(data);
let count = 0;
for (const filename of Object.keys(zip.files)) {
if (filename.includes('/')) { continue; }
const fileBuffer = await zip.files[filename].async('nodebuffer');
fs.writeFileSync(path.join(PATH_TO_DOC, filename), fileBuffer);
count++;
}
console.log(
`Successfully extracted ${count} template files to ${PATH_TO_DOC}`
);
}
const content = fs.readFileSync(PATH_TO_CSV, {encoding: 'utf8'});
const records = parse(content, {columns: true}) as FormulaRec[];
// let's group by doc id to save on document loading time
records.sort((a, b) => a.doc_id.localeCompare(b.doc_id));
if (!process.env.VERBOSE) {
log.transports.file.level = 'error'; // Suppress most of log output.
}
const docTools = createDocTools();
const session = docTools.createFakeSession('owners');
await docTools.before();
let successCount = 0;
let caseCount = 0;
fs.mkdirSync(path.join(PATH_TO_RESULTS), {recursive: true});
console.log('Testing AI assistance: ');
try {
DEPS.fetch = fetchWithCache;
let activeDoc: ActiveDoc|undefined;
for (const rec of records) {
let success: boolean = false;
let suggestedActions: AssistanceResponse['suggestedActions'] | undefined;
let newValues: CellValue[] | undefined;
let formula: string | undefined;
let history: AssistanceState = {messages: []};
let lastFollowUp: string | undefined;
// load new document
if (!activeDoc || activeDoc.docName !== rec.doc_id) {
const docPath = path.join(PATH_TO_DOC, rec.doc_id + '.grist');
activeDoc = await docTools.loadLocalDoc(docPath);
await activeDoc.waitForInitialization();
}
// get values
await activeDoc.docData!.fetchTable(rec.table_id);
const expected = activeDoc.docData!.getTable(rec.table_id)!.getColValues(rec.col_id)!.slice();
async function sendMessage(followUp?: string, rowId?: number) {
if (!activeDoc) {
throw new Error("No doc");
}
// send prompt
const tableId = rec.table_id;
const colId = rec.col_id;
const description = rec.Description;
const colInfo = await activeDoc.docStorage.get(`
select *
from _grist_Tables_column as c
left join _grist_Tables as t on t.id = c.parentId
where c.colId = ?
and t.tableId = ?
`, rec.col_id, rec.table_id);
formula = colInfo?.formula;
const result = await sendForCompletion(session, activeDoc, {
conversationId: 'conversationId',
context: {
type: 'formula',
tableId,
colId,
evaluateCurrentFormula: Boolean(followUp) && FOLLOWUP_EVALUATE,
rowId,
},
state: history,
text: followUp || description,
});
if (result.state) {
history = result.state;
}
if (rec.no_formula == "1") {
success = result.suggestedActions.length === 0;
return;
}
suggestedActions = result.suggestedActions;
if (!suggestedActions.length) {
success = false;
return;
}
// apply modification
const {actionNum} = await activeDoc.applyUserActions(session, suggestedActions);
// get new values
newValues = activeDoc.docData!.getTable(rec.table_id)!.getColValues(rec.col_id)!.slice();
// compare values
success = isEqual(expected, newValues);
if (!success && SIMULATE_CONVERSATION) {
for (let i = 0; i < expected.length; i++) {
const e = expected[i];
const v = newValues[i];
if (String(e) !== String(v)) {
const txt = `I got \`${v}\` where I expected \`${e}\`\n` +
'Please answer with the code block you (the assistant) just gave, ' +
'revised based on this information. Your answer must include a code ' +
'block. If you have to explain anything, do it after.\n';
const rowIds = activeDoc.docData!.getTable(rec.table_id)!.getRowIds();
const rowId = rowIds[i];
if (followUp) {
lastFollowUp = txt;
} else {
await sendMessage(txt, rowId);
}
break;
}
}
}
// revert modification
const [bundle] = await activeDoc.getActions([actionNum]);
await activeDoc.applyUserActionsById(session, [bundle!.actionNum], [bundle!.actionHash!], true);
}
try {
await sendMessage();
} catch (e) {
console.error(e);
}
console.log(` ${success ? 'Successfully' : 'Failed to'} complete formula ` +
`for column ${rec.table_id}.${rec.col_id} (doc=${rec.doc_id})`);
if (success) {
successCount++;
} else {
// TODO: log the difference between expected and actual, similar to what mocha does on
// failure.
// console.log('expected=', expected);
// console.log('actual=', newValues);
}
const suggestedFormula = suggestedActions?.length === 1 &&
suggestedActions[0][0] === 'ModifyColumn' &&
suggestedActions[0][3].formula || suggestedActions;
fs.writeFileSync(
path.join(
PATH_TO_RESULTS,
`${rec.table_id}_${rec.col_id}_` +
caseCount.toLocaleString('en', {minimumIntegerDigits: 8, useGrouping: false}) + '.json'),
JSON.stringify({
formula,
suggestedFormula, success,
expectedValues: expected,
suggestedValues: newValues,
history,
lastFollowUp,
}, null, 2));
caseCount++;
}
} finally {
await docTools.after();
log.transports.file.level = 'debug';
printStats();
DEPS.fetch = oldFetch;
console.log(
`AI Assistance completed ${successCount} successful prompt on a total of ${records.length};`
);
console.log(JSON.stringify(
{
hit: successCount,
total: records.length,
percentage: (100.0 * successCount) / Math.max(records.length, 1),
}
));
}
}
export function main() {
runCompletion().catch(console.error);
}
function printStats() {
console.log(`Ai assistance requests stats: ${_stats.callCount} calls`);
}
/**
* Implements a simple cache that read/write from filesystem.
*/
class JsonCache implements MapCache {
constructor() {
if (!fs.existsSync(PATH_TO_CACHE)) {
fs.mkdirSync(path.join(PATH_TO_CACHE), {recursive: true});
}
}
public get(key: string): any {
if (!this.has(key)) { return undefined; }
const content = JSON.parse(fs.readFileSync(this._path(key), 'utf8'));
return JSON.stringify(content.responseBody);
}
public has(key: string): boolean {
return fs.existsSync(this._path(key));
}
public set(key: string, value: any): JsonCache {
const content = {
requestBody: key,
responseBody: JSON.parse(value),
};
fs.writeFileSync(this._path(key), JSON.stringify(content));
return this;
}
public clear(): void {
throw new Error('not implemented');
}
public delete(_key: string): boolean {
throw new Error('not implemented');
}
private _path(key: string) {
return path.join(PATH_TO_CACHE, this._hash(key) + '.json');
}
private _hash(key: string) {
return crypto.createHash('md5').update(key).digest('hex');
}
}
/**
* Calls fetch and uses caching.
*/
const _cache = new JsonCache();
const _queue = new Map<string, any>();
async function fetchWithCache(rinfo: RequestInfo, init?: RequestInit): Promise<Response>
async function fetchWithCache(rinfo: any, init?: RequestInit): Promise<Response> {
const url: string = rinfo.url || rinfo.href || rinfo;
const hash = JSON.stringify({url, body: init?.body});
if (_cache.has(hash)) { return new Response(_cache.get(hash), {status: 200}); }
if (_queue.has(hash)) { return new Response(await _queue.get(hash), {status: 200}); }
_queue.set(hash, fetch(url, init));
const response = await _queue.get(hash);
_stats.callCount++;
if (response.status === 200) {
_cache.set(hash, await response.clone().text()); // response cannot be read twice, hence clone
}
return response;
}
// ts expect this function
fetchWithCache.isRedirect = fetch.isRedirect;