type Message = { role: string; content: string };
type State = { messages: Message[] };
const route = (s: State) =>
s.messages.at(-1)?.content.includes("搜索") ? "search" : "llm";
const search = (s: State): State => ({
messages: [...s.messages, { role: "tool", content: "搜索结果" }],
});
const llm = (s: State): State => ({
messages: [...s.messages, { role: "ai", content: "已处理" }],
});
class StateGraph<S> {
private nodes = new Map<string, (s: S) => S>();
private edges = new Map<string, string | ((s: S) => string)>();
addNode(name: string, fn: (s: S) => S) { this.nodes.set(name, fn); }
addEdge(from: string, to: string | ((s: S) => string)) { this.edges.set(from, to); }
compile() {
return {
invoke: (initial: S) => {
let curr = "START";
let state = initial;
while (true) {
const to = this.edges.get(curr);
if (!to) break;
const next = typeof to === "function" ? (to as any)(state) : to;
if (next === "END") break;
state = this.nodes.get(next)!(state);
curr = next;
}
return state;
}
};
}
}
const g = new StateGraph<State>();
g.addNode("search", search);
g.addNode("llm", llm);
g.addEdge("START", (s) => route(s));
g.addEdge("search", "llm");
g.addEdge("llm", "END");
const out = g.compile().invoke({ messages: [{ role: "user", content: "需要搜索" }] });