-
Notifications
You must be signed in to change notification settings - Fork 2k
/
token_buffer_memory.ts
138 lines (124 loc) Β· 4.2 KB
/
token_buffer_memory.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
import { ChatOpenAI } from "@langchain/openai";
import {
InputValues,
MemoryVariables,
OutputValues,
getInputValue,
getOutputValue,
} from "@langchain/core/memory";
import { getBufferString } from "@langchain/core/messages";
import {
BaseChatMemory,
BaseChatMemoryInput,
} from "@langchain/community/memory/chat_memory";
import { _formatIntermediateSteps } from "../../openai_functions/index.js";
/**
* Type definition for the fields required to initialize an instance of
* OpenAIAgentTokenBufferMemory.
*/
export type OpenAIAgentTokenBufferMemoryFields = BaseChatMemoryInput & {
llm: ChatOpenAI;
humanPrefix?: string;
aiPrefix?: string;
memoryKey?: string;
maxTokenLimit?: number;
returnMessages?: boolean;
outputKey?: string;
intermediateStepsKey?: string;
};
/**
* Memory used to save agent output and intermediate steps.
*/
export class OpenAIAgentTokenBufferMemory extends BaseChatMemory {
humanPrefix = "Human";
aiPrefix = "AI";
llm: ChatOpenAI;
memoryKey = "history";
maxTokenLimit = 12000;
returnMessages = true;
outputKey = "output";
intermediateStepsKey = "intermediateSteps";
constructor(fields: OpenAIAgentTokenBufferMemoryFields) {
super(fields);
this.humanPrefix = fields.humanPrefix ?? this.humanPrefix;
this.aiPrefix = fields.aiPrefix ?? this.aiPrefix;
this.llm = fields.llm;
this.memoryKey = fields.memoryKey ?? this.memoryKey;
this.maxTokenLimit = fields.maxTokenLimit ?? this.maxTokenLimit;
this.returnMessages = fields.returnMessages ?? this.returnMessages;
this.outputKey = fields.outputKey ?? this.outputKey;
this.intermediateStepsKey =
fields.intermediateStepsKey ?? this.intermediateStepsKey;
}
get memoryKeys(): string[] {
return [this.memoryKey];
}
/**
* Retrieves the messages from the chat history.
* @returns Promise that resolves with the messages from the chat history.
*/
async getMessages() {
return this.chatHistory.getMessages();
}
/**
* Loads memory variables from the input values.
* @param _values Input values.
* @returns Promise that resolves with the loaded memory variables.
*/
async loadMemoryVariables(_values: InputValues): Promise<MemoryVariables> {
const buffer = await this.getMessages();
if (this.returnMessages) {
return { [this.memoryKey]: buffer };
} else {
const bufferString = getBufferString(
buffer,
this.humanPrefix,
this.aiPrefix
);
return { [this.memoryKey]: bufferString };
}
}
/**
* Saves the context of the chat, including user input, AI output, and
* intermediate steps. Prunes the chat history if the total token count
* exceeds the maximum limit.
* @param inputValues Input values.
* @param outputValues Output values.
* @returns Promise that resolves when the context has been saved.
*/
async saveContext(
inputValues: InputValues,
outputValues: OutputValues
): Promise<void> {
const inputValue = getInputValue(inputValues, this.inputKey);
const outputValue = getOutputValue(outputValues, this.outputKey);
await this.chatHistory.addUserMessage(inputValue);
const intermediateStepMessages = _formatIntermediateSteps(
outputValues[this.intermediateStepsKey]
);
for (const message of intermediateStepMessages) {
await this.chatHistory.addMessage(message);
}
await this.chatHistory.addAIChatMessage(outputValue);
const currentMessages = await this.chatHistory.getMessages();
let tokenInfo = await this.llm.getNumTokensFromMessages(currentMessages);
if (tokenInfo.totalCount > this.maxTokenLimit) {
const prunedMemory = [];
while (tokenInfo.totalCount > this.maxTokenLimit) {
const retainedMessage = currentMessages.pop();
if (!retainedMessage) {
console.warn(
`Could not prune enough messages from chat history to stay under ${this.maxTokenLimit} tokens.`
);
break;
}
prunedMemory.push(retainedMessage);
tokenInfo = await this.llm.getNumTokensFromMessages(currentMessages);
}
await this.chatHistory.clear();
for (const message of prunedMemory) {
await this.chatHistory.addMessage(message);
}
}
}
}