-
Notifications
You must be signed in to change notification settings - Fork 2k
/
aws_sfn.ts
124 lines (112 loc) Β· 3.77 KB
/
aws_sfn.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
import { Tool } from "../../tools/base.js";
import {
SfnConfig,
StartExecutionAWSSfnTool,
DescribeExecutionAWSSfnTool,
SendTaskSuccessAWSSfnTool,
} from "../../tools/aws_sfn.js";
import { Toolkit } from "./base.js";
import { BaseLanguageModel } from "../../base_language/index.js";
import { renderTemplate } from "../../prompts/template.js";
import { LLMChain } from "../../chains/llm_chain.js";
import { ZeroShotAgent, ZeroShotCreatePromptArgs } from "../mrkl/index.js";
import { AgentExecutor } from "../executor.js";
export const SFN_PREFIX = `You are an agent designed to interact with AWS Step Functions state machines to execute and coordinate asynchronous workflows and tasks.
Given an input question, command, or task use the appropriate tool to execute a command to interact with AWS Step Functions and return the result.
You have access to tools for interacting with AWS Step Functions.
Given an input question, command, or task use the correct tool to complete the task.
Only use the below tools. Only use the information returned by the below tools to construct your final answer.
If the question does not seem related to AWS Step Functions or an existing state machine, just return "I don't know" as the answer.`;
export const SFN_SUFFIX = `Begin!
Question: {input}
Thought: I should look at state machines within AWS Step Functions to see what actions I can perform.
{agent_scratchpad}`;
export interface AWSSfnCreatePromptArgs extends ZeroShotCreatePromptArgs {}
/**
* Interface for the arguments required to create an AWS Step Functions
* toolkit.
*/
export interface AWSSfnToolkitArgs {
name: string;
description: string;
stateMachineArn: string;
asl?: string;
llm?: BaseLanguageModel;
}
/**
* Class representing a toolkit for interacting with AWS Step Functions.
* It initializes the AWS Step Functions tools and provides them as tools
* for the agent.
*/
export class AWSSfnToolkit extends Toolkit {
tools: Tool[];
stateMachineArn: string;
asl: string;
constructor(args: AWSSfnToolkitArgs & SfnConfig) {
super();
this.stateMachineArn = args.stateMachineArn;
if (args.asl) {
this.asl = args.asl;
}
this.tools = [
new StartExecutionAWSSfnTool({
name: args.name,
description: StartExecutionAWSSfnTool.formatDescription(
args.name,
args.description
),
stateMachineArn: args.stateMachineArn,
}),
new DescribeExecutionAWSSfnTool(
Object.assign(
args.region ? { region: args.region } : {},
args.accessKeyId && args.secretAccessKey
? {
accessKeyId: args.accessKeyId,
secretAccessKey: args.secretAccessKey,
}
: {}
)
),
new SendTaskSuccessAWSSfnTool(
Object.assign(
args.region ? { region: args.region } : {},
args.accessKeyId && args.secretAccessKey
? {
accessKeyId: args.accessKeyId,
secretAccessKey: args.secretAccessKey,
}
: {}
)
),
];
}
}
export function createAWSSfnAgent(
llm: BaseLanguageModel,
toolkit: AWSSfnToolkit,
args?: AWSSfnCreatePromptArgs
) {
const {
prefix = SFN_PREFIX,
suffix = SFN_SUFFIX,
inputVariables = ["input", "agent_scratchpad"],
} = args ?? {};
const { tools } = toolkit;
const formattedPrefix = renderTemplate(prefix, "f-string", {});
const prompt = ZeroShotAgent.createPrompt(tools, {
prefix: formattedPrefix,
suffix,
inputVariables,
});
const chain = new LLMChain({ prompt, llm });
const agent = new ZeroShotAgent({
llmChain: chain,
allowedTools: tools.map((t) => t.name),
});
return AgentExecutor.fromAgentAndTools({
agent,
tools,
returnIntermediateSteps: true,
});
}