-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
tagging.ts
91 lines (84 loc) · 3.01 KB
/
tagging.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
import { z } from "zod";
import { zodToJsonSchema } from "zod-to-json-schema";
import { JsonSchema7ObjectType } from "zod-to-json-schema/src/parsers/object.js";
import { PromptTemplate } from "../../prompts/prompt.js";
import {
FunctionParameters,
JsonOutputFunctionsParser,
} from "../../output_parsers/openai_functions.js";
import { LLMChain, LLMChainInput } from "../llm_chain.js";
import { BaseChatModel } from "../../chat_models/base.js";
import { BaseFunctionCallOptions } from "../../base_language/index.js";
/**
* Type representing the options for creating a tagging chain.
*/
export type TaggingChainOptions = {
prompt?: PromptTemplate;
} & Omit<LLMChainInput<object>, "prompt" | "llm">;
/**
* Function that returns an array of tagging functions. These functions
* are used to extract relevant information from a passage.
* @param schema The schema defining the structure of function parameters.
* @returns An array of tagging functions.
*/
function getTaggingFunctions(schema: FunctionParameters) {
return [
{
name: "information_extraction",
description: "Extracts the relevant information from the passage.",
parameters: schema,
},
];
}
const TAGGING_TEMPLATE = `Extract the desired information from the following passage.
Passage:
{input}
`;
/**
* Function that creates a tagging chain using the provided schema,
* LLM, and options. It constructs the LLM with the necessary
* functions, prompt, output parser, and tags.
* @param schema The schema defining the structure of function parameters.
* @param llm LLM to use in the chain. Must support function calling.
* @param options Options for creating the tagging chain.
* @returns A new instance of LLMChain configured for tagging.
*/
export function createTaggingChain(
schema: FunctionParameters,
llm: BaseChatModel<BaseFunctionCallOptions>,
options: TaggingChainOptions = {}
) {
const { prompt = PromptTemplate.fromTemplate(TAGGING_TEMPLATE), ...rest } =
options;
const functions = getTaggingFunctions(schema);
const outputParser = new JsonOutputFunctionsParser();
return new LLMChain({
llm,
prompt,
llmKwargs: { functions },
outputParser,
tags: ["openai_functions", "tagging"],
...rest,
});
}
/**
* Function that creates a tagging chain from a Zod schema. It converts
* the Zod schema to a JSON schema using the zodToJsonSchema function and
* then calls createTaggingChain with the converted schema.
* @param schema The Zod schema which extracted data should match.
* @param llm LLM to use in the chain. Must support function calling.
* @param options Options for creating the tagging chain.
* @returns A new instance of LLMChain configured for tagging.
*/
export function createTaggingChainFromZod(
// eslint-disable-next-line @typescript-eslint/no-explicit-any
schema: z.ZodObject<any, any, any, any>,
llm: BaseChatModel<BaseFunctionCallOptions>,
options?: TaggingChainOptions
) {
return createTaggingChain(
zodToJsonSchema(schema) as JsonSchema7ObjectType,
llm,
options
);
}