gristlabs_grist-core/test/formula-dataset/runCompletion_impl.ts

360 lines
12 KiB
TypeScript
Raw Normal View History

/**
* 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/lib/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 {
(core) Modify prompt so that model may say it cannot help with certain requests. Summary: This tweaks the prompting so that the user's message is given on its own instead of as a docstring within Python. This is so that the prompt makes sense when: - the user asks a question such as "Can you write me a formula which does ...?" rather than describing their formula as a docstring would, or - the user sends a message that doesn't ask for a formula at all (https://grist.slack.com/archives/C0234CPPXPA/p1687699944315069?thread_ts=1687698078.832209&cid=C0234CPPXPA) Also added wording for the model to refuse when the user asks for something that the model cannot do. Because the code (and maybe in some cases the model) for non-ChatGPT models relies on the prompt consisting entirely of Python code produced by the data engine (which no longer contains the user's message) those code paths have been disabled for now. Updating them now seems like undesirable drag, I think it'd be better to revisit this when iteration/experimentation has slowed down and stabilised. Test Plan: Added entries to the formula dataset where the response shouldn't contain a formula, indicated by the value `1` for the new column `no_formula`. This is somewhat successful, as the model does refuse to help in some of the new test cases, but not all. Performance on existing entries also seems a bit worse, but it's hard to distinguish this from random noise. Hopefully this can be remedied in the future with more work, e.g. automatic followup messages containing example inputs and outputs. Reviewers: paulfitz Reviewed By: paulfitz Subscribers: dsagal Differential Revision: https://phab.getgrist.com/D3936
2023-06-27 11:39:15 +00:00
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;
}
(core) Modify prompt so that model may say it cannot help with certain requests. Summary: This tweaks the prompting so that the user's message is given on its own instead of as a docstring within Python. This is so that the prompt makes sense when: - the user asks a question such as "Can you write me a formula which does ...?" rather than describing their formula as a docstring would, or - the user sends a message that doesn't ask for a formula at all (https://grist.slack.com/archives/C0234CPPXPA/p1687699944315069?thread_ts=1687698078.832209&cid=C0234CPPXPA) Also added wording for the model to refuse when the user asks for something that the model cannot do. Because the code (and maybe in some cases the model) for non-ChatGPT models relies on the prompt consisting entirely of Python code produced by the data engine (which no longer contains the user's message) those code paths have been disabled for now. Updating them now seems like undesirable drag, I think it'd be better to revisit this when iteration/experimentation has slowed down and stabilised. Test Plan: Added entries to the formula dataset where the response shouldn't contain a formula, indicated by the value `1` for the new column `no_formula`. This is somewhat successful, as the model does refuse to help in some of the new test cases, but not all. Performance on existing entries also seems a bit worse, but it's hard to distinguish this from random noise. Hopefully this can be remedied in the future with more work, e.g. automatic followup messages containing example inputs and outputs. Reviewers: paulfitz Reviewed By: paulfitz Subscribers: dsagal Differential Revision: https://phab.getgrist.com/D3936
2023-06-27 11:39:15 +00:00
if (rec.no_formula == "1") {
success = result.suggestedActions.length === 0;
return;
(core) Modify prompt so that model may say it cannot help with certain requests. Summary: This tweaks the prompting so that the user's message is given on its own instead of as a docstring within Python. This is so that the prompt makes sense when: - the user asks a question such as "Can you write me a formula which does ...?" rather than describing their formula as a docstring would, or - the user sends a message that doesn't ask for a formula at all (https://grist.slack.com/archives/C0234CPPXPA/p1687699944315069?thread_ts=1687698078.832209&cid=C0234CPPXPA) Also added wording for the model to refuse when the user asks for something that the model cannot do. Because the code (and maybe in some cases the model) for non-ChatGPT models relies on the prompt consisting entirely of Python code produced by the data engine (which no longer contains the user's message) those code paths have been disabled for now. Updating them now seems like undesirable drag, I think it'd be better to revisit this when iteration/experimentation has slowed down and stabilised. Test Plan: Added entries to the formula dataset where the response shouldn't contain a formula, indicated by the value `1` for the new column `no_formula`. This is somewhat successful, as the model does refuse to help in some of the new test cases, but not all. Performance on existing entries also seems a bit worse, but it's hard to distinguish this from random noise. Hopefully this can be remedied in the future with more work, e.g. automatic followup messages containing example inputs and outputs. Reviewers: paulfitz Reviewed By: paulfitz Subscribers: dsagal Differential Revision: https://phab.getgrist.com/D3936
2023-06-27 11:39:15 +00:00
}
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;