@langgraph-js/sdk
Version:
The UI SDK for LangGraph - seamlessly integrate your AI agents with frontend interfaces
310 lines (293 loc) • 10.8 kB
text/typescript
import { RenderMessage } from "./LangGraphClient.js";
import type { Message } from "@langchain/langgraph-sdk";
import { CallToolResult, UnionTool } from "./tool/createTool.js";
import { ToolRenderData } from "./tool/ToolUI.js";
import { createChatStore } from "./ui-store/createChatStore.js";
/**
* @zh 测试任务接口
* @en Test task interface
*/
interface TestTask {
/** 任务是否成功完成 */
success: boolean;
/** 执行任务的函数 */
runTask: (messages: readonly RenderMessage[]) => Promise<void>;
/** 任务失败时的回调函数 */
fail: () => void;
}
export class TestLogger {
info(message: string) {
console.log(message);
}
logMessage(message: RenderMessage) {
const emoji = message.type === "ai" ? "🤖" : message.type === "human" ? "👤" : "🔧";
const header = `${emoji} ${message.type} | ${message.name ?? "null"} | ${message.id}`;
if (message.type === "tool") {
return console.log(`${header}
🔧 Input: ${message.tool_input?.slice(0, 100) ?? ""}
💬 Output: ${message.content?.slice(0, 100) ?? ""}
`);
}
console.log(`---
${header}
💬 Output: ${message.content?.slice(0, 100) ?? ""}
`);
}
}
/**
* @zh LangGraph 测试工具，可以配合 vitest 等常用框架进行测试
* @en LangGraph test tool, can be used with vitest and other common frameworks for testing
*
* @example
* ```typescript
* const testChat = new TestLangGraphChat(createLangGraphClient(), { debug: true });
* await testChat.humanInput("你好", async () => {
* const aiMessage = await testChat.waitFor("ai");
* expect(aiMessage.content).toBeDefined();
* });
* ```
*/
export class TestLangGraphChat {
/** 是否开启调试模式 */
private debug = false;
/** 上次消息数量，用于检测消息变化 */
private lastLength = 0;
/** 待处理的测试任务列表 */
protected processFunc: TestTask[] = [];
/** 自定义日志器 */
private logger: TestLogger;
/**
* @zh 构造函数，初始化测试环境
* @en Constructor, initialize test environment
*/
constructor(
readonly store: ReturnType<typeof createChatStore>,
options: {
debug?: boolean;
tools?: UnionTool<any, any, any>[];
logger?: TestLogger;
}
) {
this.debug = options.debug ?? false;
this.logger = options.logger ?? new TestLogger();
options.tools && this.addTools(options.tools);
const renderMessages = this.store.data.renderMessages;
// 订阅消息变化，自动检查任务完成状态
renderMessages.subscribe((messages) => {
this.checkAllTask(messages);
});
}
/**
* @zh 获取当前所有渲染消息
* @en Get all current render messages
*/
getMessages() {
return this.store.data.renderMessages.get();
}
/**
* @zh 添加工具到测试环境中，会自动包装工具的 execute 方法
* @en Add tools to test environment, automatically wraps tool execute methods
*
* @example
* ```typescript
* const tools = [createUITool({ name: "test_tool", ... })];
* testChat.addTools(tools);
* ```
*/
addTools(tools: UnionTool<any, any, any>[]) {
tools.forEach((tool) => {
if (tool.execute) {
const oldExecute = tool.execute;
// 包装原始的 execute 方法，在执行后触发任务检查
tool.execute = (...args) => {
setTimeout(() => {
this.checkAllTask(this.getMessages(), {
skipLengthCheck: true,
});
}, 10);
return oldExecute!(...args);
};
}
});
this.store.mutations.setTools(tools);
}
/**
* @zh 设置额外参数
* @en Set extra states to LangGraph
*
* @example
* ```typescript
* testChat.setExtraParams({
* extraParam: "value",
* });
*/
setExtraParams(extraParams: Record<string, any>) {
const client = this.store.data.client.get();
if (client) {
client.extraParams = extraParams;
}
}
/**
* @zh 检查所有待处理的测试任务，只有在消息数量发生变化时才执行检查
* @en Check all pending test tasks, only executes when message count changes
*/
checkAllTask(messages: readonly RenderMessage[], options: { skipLengthCheck?: boolean } = {}) {
// 只有 lastLength 发生变化时，才执行检查
if (!options.skipLengthCheck && this.lastLength === messages.length) {
return;
}
this.lastLength = messages.length;
// 执行所有未完成的任务
for (const task of this.processFunc) {
!task.success && task.runTask(options.skipLengthCheck ? messages : messages.slice(0, -1));
}
// 调试模式下打印最新消息
const item = messages[messages.length - (options.skipLengthCheck ? 1 : 2)];
if (this.debug && item) {
this.logger.logMessage(item);
}
}
private readited = false;
/**
* @zh 准备测试环境，初始化客户端连接
* @en Prepare test environment, initialize client connection
*/
ready() {
if (this.readited) {
return;
}
this.readited = true;
return this.store.mutations.initClient();
}
/**
* @zh 模拟人类输入消息并等待测试任务完成，这是测试的核心方法
* @en Simulate human input and wait for test tasks to complete, this is the core test method
*
* @example
* ```typescript
* await testChat.humanInput("请帮我思考一下", async () => {
* const toolMessage = await testChat.waitFor("tool", "thinking");
* expect(toolMessage.tool_input).toBeDefined();
*
* const aiMessage = await testChat.waitFor("ai");
* expect(aiMessage.content).toContain("思考");
* });
* ```
*/
async humanInput(text: Message["content"], context: () => Promise<void>) {
await this.ready();
// console.log(text);
return Promise.all([
context(),
this.store.mutations
.sendMessage([
{
type: "human",
content: text,
},
])
.then(async () => {
// messages 有 10 ms 的 debounce，我们需要稍等一下
await new Promise((resolve) => setTimeout(resolve, 100));
this.checkAllTask(this.getMessages(), {
skipLengthCheck: true,
});
})
.then(async (res) => {
// 检查是否还有未完成的任务
const tasks = this.processFunc.filter((i) => {
return !i.success;
});
if (tasks.length) {
console.warn("still have ", tasks.length, " tasks");
await Promise.all(tasks.map((i) => i.fail()));
throw new Error("test task failed");
}
this.processFunc = [];
return res;
}),
]);
}
/**
* @zh 等待特定类型的消息出现，创建异步等待任务
* @en Wait for specific type of message to appear, creates async waiting task
*
* @example
* ```typescript
* // 等待 AI 回复
* const aiMessage = await testChat.waitFor("ai");
*
* // 等待特定工具调用
* const toolMessage = await testChat.waitFor("tool", "sequential-thinking");
* ```
*/
waitFor<D extends "tool" | "ai", T extends RenderMessage, N extends D extends "tool" ? string : undefined>(type: D, name?: N): Promise<T> {
return new Promise((resolve, reject) => {
this.processFunc.push({
success: false,
async runTask(messages) {
const lastMessage = messages[messages.length - 1];
if (!lastMessage) {
return;
}
// 检查消息类型和名称是否匹配
if (lastMessage.type === type && (name ? lastMessage.name === name : true)) {
resolve(lastMessage as T);
this.success = true;
}
},
fail() {
reject(new Error(`wait for ${type} ${name} failed`));
},
});
});
}
/**
* @zh 响应前端工具调用，模拟用户对工具的响应
* @en Respond to frontend tool calls, simulates user response to tools
*
* @example
* ```typescript
* const toolMessage = await testChat.waitFor("tool", "ask_user_for_approve");
* await testChat.responseFeTool(toolMessage, "approved");
* ```
*/
async responseFeTool(message: RenderMessage, value: CallToolResult) {
if (message.content) {
throw new Error(`message is Done. content: ${message.content}`);
}
const tool = new ToolRenderData(message, this.store.data.client.get()!);
tool.response(value);
const messages = await this.waitFor("tool", message.name!);
if (messages.content) {
return messages;
}
throw new Error("tool response failed");
}
/**
* @zh 查找最后一条指定类型的消息，从消息数组末尾开始向前查找
* @en Find the last message of specified type, searches backwards from end of messages
*
* @example
* ```typescript
* // 查找最后一条 AI 消息
* const lastAI = testChat.findLast("ai");
*
* // 查找最后一条人类消息
* const lastHuman = testChat.findLast("human");
* ```
*/
findLast(type: "human" | "ai" | "tool", options: { before?: (item: RenderMessage) => boolean } = {}) {
const messages = this.getMessages();
for (let i = messages.length - 1; i >= 0; i--) {
const item = messages[i];
if (type === item.type) {
return item;
}
if (options.before && options.before(item)) {
throw new Error(`${type} not found; before specified`);
}
}
throw new Error(`${type} not found `);
}
}