-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
runner_utils.ts
519 lines (489 loc) Β· 15.9 KB
/
runner_utils.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
import { BaseLanguageModel } from "@langchain/core/language_models/base";
import { Serialized } from "@langchain/core/load/serializable";
import { mapStoredMessagesToChatMessages } from "@langchain/core/messages";
import {
Runnable,
RunnableConfig,
RunnableLambda,
} from "@langchain/core/runnables";
import { RunCollectorCallbackHandler } from "@langchain/core/tracers/run_collector";
import { LangChainTracer } from "@langchain/core/tracers/tracer_langchain";
import { ChainValues } from "@langchain/core/utils/types";
import { Client, Example, Feedback, Run } from "langsmith";
import { EvaluationResult, RunEvaluator } from "langsmith/evaluation";
import { DataType } from "langsmith/schemas";
import { LLMStringEvaluator } from "../evaluation/base.js";
import { loadEvaluator } from "../evaluation/loader.js";
import { EvaluatorType } from "../evaluation/types.js";
import type {
DynamicRunEvaluatorParams,
EvalConfig,
EvaluatorInputFormatter,
RunEvalConfig,
RunEvaluatorLike,
} from "./config.js";
import { randomName } from "./name_generation.js";
import { ProgressBar } from "./progress.js";
export type ChainOrFactory =
| Runnable
| (() => Runnable)
// eslint-disable-next-line @typescript-eslint/no-explicit-any
| ((obj: any) => any)
// eslint-disable-next-line @typescript-eslint/no-explicit-any
| ((obj: any) => Promise<any>)
| (() => (obj: unknown) => unknown)
| (() => (obj: unknown) => Promise<unknown>);
class RunIdExtractor {
runIdPromiseResolver: (runId: string) => void;
runIdPromise: Promise<string>;
constructor() {
this.runIdPromise = new Promise<string>((extract) => {
this.runIdPromiseResolver = extract;
});
}
handleChainStart = (
_chain: Serialized,
_inputs: ChainValues,
runId: string
) => {
this.runIdPromiseResolver(runId);
};
async extract(): Promise<string> {
return this.runIdPromise;
}
}
/**
* Wraps an evaluator function + implements the RunEvaluator interface.
*/
class DynamicRunEvaluator implements RunEvaluator {
evaluator: RunnableLambda<DynamicRunEvaluatorParams, EvaluationResult>;
constructor(evaluator: RunEvaluatorLike) {
this.evaluator = new RunnableLambda({ func: evaluator });
}
/**
* Evaluates a run with an optional example and returns the evaluation result.
* @param run The run to evaluate.
* @param example The optional example to use for evaluation.
* @returns A promise that extracts to the evaluation result.
*/
async evaluateRun(run: Run, example?: Example): Promise<EvaluationResult> {
const extractor = new RunIdExtractor();
const tracer = new LangChainTracer({ projectName: "evaluators" });
const result = await this.evaluator.invoke(
{
run,
example,
input: run.inputs,
prediction: run.outputs,
reference: example?.outputs,
},
{
callbacks: [extractor, tracer],
}
);
const runId = await extractor.extract();
return {
sourceRunId: runId,
...result,
};
}
}
// eslint-disable-next-line @typescript-eslint/no-explicit-any
function isLLMStringEvaluator(evaluator: any): evaluator is LLMStringEvaluator {
return evaluator && typeof evaluator.evaluateStrings === "function";
}
/**
* Wraps an off-the-shelf evaluator (loaded using loadEvaluator; of EvaluatorType[T])
* and composes with a prepareData function so the user can prepare the trace and
* dataset data for the evaluator.
*/
class PreparedRunEvaluator implements RunEvaluator {
evaluator: LLMStringEvaluator;
formatEvaluatorInputs: EvaluatorInputFormatter;
isStringEvaluator: boolean;
evaluationName: string;
constructor(
evaluator: LLMStringEvaluator,
evaluationName: string,
formatEvaluatorInputs: EvaluatorInputFormatter
) {
this.evaluator = evaluator;
this.isStringEvaluator = typeof evaluator?.evaluateStrings === "function";
this.evaluationName = evaluationName;
this.formatEvaluatorInputs = formatEvaluatorInputs;
}
static async fromEvalConfig(
config: EvalConfig | keyof EvaluatorType
): Promise<PreparedRunEvaluator> {
const evaluatorType =
typeof config === "string" ? config : config.evaluatorType;
const evalConfig = typeof config === "string" ? ({} as EvalConfig) : config;
const evaluator = await loadEvaluator(evaluatorType, evalConfig);
const feedbackKey = evalConfig?.feedbackKey ?? evaluator?.evaluationName;
if (!feedbackKey) {
throw new Error(
`Evaluator of type ${evaluatorType} must have an evaluationName` +
` or feedbackKey. Please manually provide a feedbackKey in the EvalConfig.`
);
}
if (!isLLMStringEvaluator(evaluator)) {
throw new Error(
`Evaluator of type ${evaluatorType} not yet supported. ` +
"Please use a string evaluator, or implement your " +
"evaluation logic as a customEvaluator."
);
}
return new PreparedRunEvaluator(
evaluator as LLMStringEvaluator,
feedbackKey,
evalConfig?.formatEvaluatorInputs
);
}
/**
* Evaluates a run with an optional example and returns the evaluation result.
* @param run The run to evaluate.
* @param example The optional example to use for evaluation.
* @returns A promise that extracts to the evaluation result.
*/
async evaluateRun(run: Run, example?: Example): Promise<EvaluationResult> {
const { prediction, input, reference } = this.formatEvaluatorInputs({
rawInput: run.inputs,
rawPrediction: run.outputs,
rawReferenceOutput: example?.outputs,
run,
});
const extractor = new RunIdExtractor();
const tracer = new LangChainTracer({ projectName: "evaluators" });
if (this.isStringEvaluator) {
const evalResult = await this.evaluator.evaluateStrings(
{
prediction: prediction as string,
reference: reference as string,
input: input as string,
},
{
callbacks: [extractor, tracer],
}
);
const runId = await extractor.extract();
return {
key: this.evaluationName,
comment: evalResult?.reasoning,
sourceRunId: runId,
...evalResult,
};
}
throw new Error(
"Evaluator not yet supported. " +
"Please use a string evaluator, or implement your " +
"evaluation logic as a customEvaluator."
);
}
}
class LoadedEvalConfig {
constructor(public evaluators: (RunEvaluator | DynamicRunEvaluator)[]) {}
static async fromRunEvalConfig(
config: RunEvalConfig
): Promise<LoadedEvalConfig> {
// Custom evaluators are applied "as-is"
const customEvaluators = config?.customEvaluators?.map((evaluator) => {
if (typeof evaluator === "function") {
return new DynamicRunEvaluator(evaluator);
} else {
return evaluator;
}
});
const offTheShelfEvaluators = await Promise.all(
config?.evaluators?.map(
async (evaluator) =>
await PreparedRunEvaluator.fromEvalConfig(evaluator)
) ?? []
);
return new LoadedEvalConfig(
(customEvaluators ?? []).concat(offTheShelfEvaluators ?? [])
);
}
}
export type RunOnDatasetParams = {
evaluationConfig?: RunEvalConfig;
projectMetadata?: Record<string, unknown>;
projectName?: string;
client?: Client;
maxConcurrency?: number;
};
/**
* Internals expect a constructor () -> Runnable. This function wraps/coerces
* the provided LangChain object, custom function, or factory function into
* a constructor of a runnable.
* @param modelOrFactory The model or factory to create a wrapped model from.
* @returns A function that returns the wrapped model.
* @throws Error if the modelOrFactory is invalid.
*/
const createWrappedModel = async (modelOrFactory: ChainOrFactory) => {
if (Runnable.isRunnable(modelOrFactory)) {
return () => modelOrFactory;
}
if (typeof modelOrFactory === "function") {
try {
// If it works with no arguments, assume it's a factory
let res = (modelOrFactory as () => Runnable)();
if (
res &&
typeof (res as unknown as Promise<Runnable>).then === "function"
) {
res = await res;
}
return modelOrFactory as () => Runnable;
} catch (err) {
// Otherwise, it's a custom UDF, and we'll wrap
// in a lambda
const wrappedModel = new RunnableLambda({ func: modelOrFactory });
return () => wrappedModel;
}
}
throw new Error("Invalid modelOrFactory");
};
const loadExamples = async ({
datasetName,
client,
projectName,
}: {
datasetName: string;
client: Client;
projectName: string;
maxConcurrency: number;
}) => {
const exampleIterator = client.listExamples({ datasetName });
const configs: RunnableConfig[] = [];
const runCollectors = [];
const examples = [];
for await (const example of exampleIterator) {
const runCollector = new RunCollectorCallbackHandler({
exampleId: example.id,
});
configs.push({
callbacks: [
new LangChainTracer({ exampleId: example.id, projectName }),
runCollector,
],
});
examples.push(example);
runCollectors.push(runCollector);
}
return {
configs,
examples,
runCollectors,
};
};
const applyEvaluators = async ({
evaluation,
runs,
examples,
client,
}: {
evaluation: LoadedEvalConfig;
runs: Run[];
examples: Example[];
client: Client;
}) => {
// TODO: Parallelize and/or put in callbacks to speed up evals.
const { evaluators } = evaluation;
const progress = new ProgressBar({
total: examples.length,
format: "Running Evaluators: {bar} {percentage}% | {value}/{total}\n",
});
const results: Record<
string,
{ run_id: string; execution_time?: number; feedback: Feedback[] }
> = {};
for (let i = 0; i < runs.length; i += 1) {
const run = runs[i];
const example = examples[i];
const evaluatorResults = await Promise.all(
evaluators.map((evaluator) =>
client.evaluateRun(run, evaluator, {
referenceExample: example,
loadChildRuns: false,
})
)
);
progress.increment();
results[example.id] = {
execution_time:
run?.end_time && run.start_time
? run.end_time - run.start_time
: undefined,
feedback: evaluatorResults,
run_id: run.id,
};
}
return results;
};
export type EvalResults = {
projectName: string;
results: {
[key: string]: {
execution_time?: number;
run_id: string;
feedback: Feedback[];
};
};
};
const getExamplesInputs = (
examples: Example[],
chainOrFactory: ChainOrFactory,
dataType?: DataType
) => {
if (dataType === "chat") {
// For some batty reason, we store the chat dataset differently.
// { type: "system", data: { content: inputs.input } },
// But we need to create AIMesage, SystemMessage, etc.
return examples.map(({ inputs }) =>
mapStoredMessagesToChatMessages(inputs.input)
);
}
// If it's a language model and ALL example inputs have a single value,
// then we can be friendly and flatten the inputs to a list of strings.
const isLanguageModel =
typeof chainOrFactory === "object" &&
typeof (chainOrFactory as BaseLanguageModel)._llmType === "function";
if (
isLanguageModel &&
examples.every(({ inputs }) => Object.keys(inputs).length === 1)
) {
return examples.map(({ inputs }) => Object.values(inputs)[0]);
}
return examples.map(({ inputs }) => inputs);
};
/**
* Evaluates a given model or chain against a specified LangSmith dataset.
*
* This function fetches example records from the specified dataset,
* runs the model or chain against each example, and returns the evaluation
* results.
*
* @param chainOrFactory - A model or factory/constructor function to be evaluated. It can be a
* Runnable instance, a factory function that returns a Runnable, or a user-defined
* function or factory.
*
* @param datasetName - The name of the dataset against which the evaluation will be
* performed. This dataset should already be defined and contain the relevant data
* for evaluation.
*
* @param options - (Optional) Additional parameters for the evaluation process:
* - `evaluationConfig` (RunEvalConfig): Configuration for the evaluation, including
* standard and custom evaluators.
* - `projectName` (string): Name of the project for logging and tracking.
* - `projectMetadata` (Record<string, unknown>): Additional metadata for the project.
* - `client` (Client): Client instance for LangChain service interaction.
* - `maxConcurrency` (number): Maximum concurrency level for dataset processing.
*
* @returns A promise that resolves to an `EvalResults` object. This object includes
* detailed results of the evaluation, such as execution time, run IDs, and feedback
* for each entry in the dataset.
*
* @example
* ```typescript
* // Example usage for evaluating a model on a dataset
* async function evaluateModel() {
* const chain = /* ...create your model or chain...*\//
* const datasetName = 'example-dataset';
* const client = new Client(/* ...config... *\//);
*
* const evaluationConfig = {
* evaluators: [/* ...evaluators... *\//],
* customEvaluators: [/* ...custom evaluators... *\//],
* };
*
* const results = await runOnDataset(chain, datasetName, {
* evaluationConfig,
* client,
* });
*
* console.log('Evaluation Results:', results);
* }
*
* evaluateModel();
* ```
* In this example, `runOnDataset` is used to evaluate a language model (or a chain of models) against
* a dataset named 'example-dataset'. The evaluation process is configured using `RunEvalConfig`, which can
* include both standard and custom evaluators. The `Client` instance is used to interact with LangChain services.
* The function returns the evaluation results, which can be logged or further processed as needed.
*/
export const runOnDataset = async (
chainOrFactory: ChainOrFactory,
datasetName: string,
{
evaluationConfig,
projectName,
projectMetadata,
client,
maxConcurrency,
}: RunOnDatasetParams
) => {
const wrappedModel = await createWrappedModel(chainOrFactory);
const testClient = client ?? new Client();
const testProjectName = projectName ?? randomName();
const dataset = await testClient.readDataset({ datasetName });
const datasetId = dataset.id;
const testConcurrency = maxConcurrency ?? 5;
const { configs, examples, runCollectors } = await loadExamples({
datasetName,
client: testClient,
projectName: testProjectName,
maxConcurrency: testConcurrency,
});
await testClient.createProject({
projectName: testProjectName,
referenceDatasetId: datasetId,
projectExtra: { metadata: { ...projectMetadata } },
});
const wrappedRunnable: Runnable = new RunnableLambda({
func: wrappedModel,
}).withConfig({ runName: "evaluationRun" });
const runInputs = getExamplesInputs(
examples,
chainOrFactory,
dataset.data_type
);
const progress = new ProgressBar({
total: runInputs.length,
format: "Predicting: {bar} {percentage}% | {value}/{total}",
});
// TODO: Collect the runs as well.
await wrappedRunnable
.withListeners({
onEnd: () => progress.increment(),
})
// TODO: Insert evaluation inline for immediate feedback.
.batch(runInputs, configs, {
maxConcurrency,
returnExceptions: true,
});
progress.complete();
const runs: Run[] = [];
for (let i = 0; i < examples.length; i += 1) {
runs.push(runCollectors[i].tracedRuns[0]);
}
let evalResults: Record<
string,
{ run_id: string; execution_time?: number; feedback: Feedback[] }
> = {};
if (evaluationConfig) {
const loadedEvalConfig = await LoadedEvalConfig.fromRunEvalConfig(
evaluationConfig
);
evalResults = await applyEvaluators({
evaluation: loadedEvalConfig,
runs,
examples,
client: testClient,
});
}
const results: EvalResults = {
projectName: testProjectName,
results: evalResults ?? {},
};
return results;
};