-
Notifications
You must be signed in to change notification settings - Fork 2k
/
extraction.ts
87 lines (82 loc) 路 2.94 KB
/
extraction.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
import { z } from "zod";
import { zodToJsonSchema, JsonSchema7ObjectType } from "zod-to-json-schema";
import { BaseChatModel } from "@langchain/core/language_models/chat_models";
import { BaseFunctionCallOptions } from "@langchain/core/language_models/base";
import { PromptTemplate } from "@langchain/core/prompts";
import {
FunctionParameters,
JsonKeyOutputFunctionsParser,
} from "../../output_parsers/openai_functions.js";
import { LLMChain } from "../llm_chain.js";
/**
* Function that returns an array of extraction functions. These functions
* are used to extract relevant information from a passage.
* @param schema The schema of the function parameters.
* @returns An array of extraction functions.
*/
function getExtractionFunctions(schema: FunctionParameters) {
return [
{
name: "information_extraction",
description: "Extracts the relevant information from the passage.",
parameters: {
type: "object",
properties: {
info: {
type: "array",
items: {
type: schema.type,
properties: schema.properties,
required: schema.required,
},
},
},
required: ["info"],
},
},
];
}
const _EXTRACTION_TEMPLATE = `Extract and save the relevant entities mentioned in the following passage together with their properties.
Passage:
{input}
`;
/**
* Function that creates an extraction chain using the provided JSON schema.
* It sets up the necessary components, such as the prompt, output parser, and tags.
* @param schema JSON schema of the function parameters.
* @param llm Must be a ChatOpenAI or AnthropicFunctions model that supports function calling.
* @returns A LLMChain instance configured to return data matching the schema.
*/
export function createExtractionChain(
schema: FunctionParameters,
llm: BaseChatModel<BaseFunctionCallOptions>
) {
const functions = getExtractionFunctions(schema);
const prompt = PromptTemplate.fromTemplate(_EXTRACTION_TEMPLATE);
const outputParser = new JsonKeyOutputFunctionsParser({ attrName: "info" });
return new LLMChain({
llm,
prompt,
llmKwargs: { functions },
outputParser,
tags: ["openai_functions", "extraction"],
});
}
/**
* Function that creates an extraction chain from a Zod schema. It
* converts the Zod schema to a JSON schema using zod-to-json-schema
* before creating the extraction chain.
* @param schema The Zod schema which extracted data should match
* @param llm Must be a ChatOpenAI or AnthropicFunctions model that supports function calling.
* @returns A LLMChain instance configured to return data matching the schema.
*/
export function createExtractionChainFromZod(
// eslint-disable-next-line @typescript-eslint/no-explicit-any
schema: z.ZodObject<any, any, any, any>,
llm: BaseChatModel<BaseFunctionCallOptions>
) {
return createExtractionChain(
zodToJsonSchema(schema) as JsonSchema7ObjectType,
llm
);
}