-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
base.ts
157 lines (142 loc) Β· 4.3 KB
/
base.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
import { z } from "zod";
import {
CallbackManager,
CallbackManagerForToolRun,
Callbacks,
parseCallbackConfigArg,
} from "../callbacks/manager.js";
import { BaseLangChain, BaseLangChainParams } from "../base_language/index.js";
import { RunnableConfig } from "../schema/runnable/config.js";
/**
* Parameters for the Tool classes.
*/
export interface ToolParams extends BaseLangChainParams {}
/**
* Custom error class used to handle exceptions related to tool input parsing.
* It extends the built-in `Error` class and adds an optional `output`
* property that can hold the output that caused the exception.
*/
export class ToolInputParsingException extends Error {
output?: string;
constructor(message: string, output?: string) {
super(message);
this.output = output;
}
}
/**
* Base class for Tools that accept input of any shape defined by a Zod schema.
*/
export abstract class StructuredTool<
// eslint-disable-next-line @typescript-eslint/no-explicit-any
T extends z.ZodObject<any, any, any, any> = z.ZodObject<any, any, any, any>
> extends BaseLangChain<
(z.output<T> extends string ? string : never) | z.input<T>,
string
> {
abstract schema: T | z.ZodEffects<T>;
get lc_namespace() {
return ["langchain", "tools"];
}
constructor(fields?: ToolParams) {
super(fields ?? {});
}
protected abstract _call(
arg: z.output<T>,
runManager?: CallbackManagerForToolRun
): Promise<string>;
/**
* Invokes the tool with the provided input and configuration.
* @param input The input for the tool.
* @param config Optional configuration for the tool.
* @returns A Promise that resolves with a string.
*/
async invoke(
input: (z.output<T> extends string ? string : never) | z.input<T>,
config?: RunnableConfig
): Promise<string> {
return this.call(input, config);
}
/**
* Calls the tool with the provided argument, configuration, and tags. It
* parses the input according to the schema, handles any errors, and
* manages callbacks.
* @param arg The input argument for the tool.
* @param configArg Optional configuration or callbacks for the tool.
* @param tags Optional tags for the tool.
* @returns A Promise that resolves with a string.
*/
async call(
arg: (z.output<T> extends string ? string : never) | z.input<T>,
configArg?: Callbacks | RunnableConfig,
/** @deprecated */
tags?: string[]
): Promise<string> {
let parsed;
try {
parsed = await this.schema.parseAsync(arg);
} catch (e) {
throw new ToolInputParsingException(
`Received tool input did not match expected schema`,
JSON.stringify(arg)
);
}
const config = parseCallbackConfigArg(configArg);
const callbackManager_ = await CallbackManager.configure(
config.callbacks,
this.callbacks,
config.tags || tags,
this.tags,
config.metadata,
this.metadata,
{ verbose: this.verbose }
);
const runManager = await callbackManager_?.handleToolStart(
this.toJSON(),
typeof parsed === "string" ? parsed : JSON.stringify(parsed),
undefined,
undefined,
undefined,
undefined,
config.runName
);
let result;
try {
result = await this._call(parsed, runManager);
} catch (e) {
await runManager?.handleToolError(e);
throw e;
}
await runManager?.handleToolEnd(result);
return result;
}
abstract name: string;
abstract description: string;
returnDirect = false;
}
/**
* Base class for Tools that accept input as a string.
*/
export abstract class Tool extends StructuredTool {
schema = z
.object({ input: z.string().optional() })
.transform((obj) => obj.input);
constructor(fields?: ToolParams) {
super(fields);
}
/**
* Calls the tool with the provided argument and callbacks. It handles
* string inputs specifically.
* @param arg The input argument for the tool, which can be a string, undefined, or an input of the tool's schema.
* @param callbacks Optional callbacks for the tool.
* @returns A Promise that resolves with a string.
*/
call(
arg: string | undefined | z.input<this["schema"]>,
callbacks?: Callbacks
): Promise<string> {
return super.call(
typeof arg === "string" || !arg ? { input: arg } : arg,
callbacks
);
}
}