import { StateGraph, Annotation, START, END } from "@langchain/langgraph";
import { createAgent, tool } from "langchain";
import { ChatOpenAI, OpenAIEmbeddings } from "@langchain/openai";
import { MemoryVectorStore } from "@langchain/classic/vectorstores/memory";
import * as z from "zod";
const State = Annotation.Root({
question: Annotation<string>(),
rewrittenQuery: Annotation<string>(),
documents: Annotation<string[]>(),
answer: Annotation<string>(),
});
// 包含阵容、比赛结果和球员统计数据的 WNBA 知识库
const embeddings = new OpenAIEmbeddings();
const vectorStore = await MemoryVectorStore.fromTexts(
[
// 阵容
"New York Liberty 2024 roster: Breanna Stewart, Sabrina Ionescu, Jonquel Jones, Courtney Vandersloot.",
"Las Vegas Aces 2024 roster: A'ja Wilson, Kelsey Plum, Jackie Young, Chelsea Gray.",
"Indiana Fever 2024 roster: Caitlin Clark, Aliyah Boston, Kelsey Mitchell, NaLyssa Smith.",
// 比赛结果
"2024 WNBA Finals: New York Liberty defeated Minnesota Lynx 3-2 to win the championship.",
"June 15, 2024: Indiana Fever 85, Chicago Sky 79. Caitlin Clark had 23 points and 8 assists.",
"August 20, 2024: Las Vegas Aces 92, Phoenix Mercury 84. A'ja Wilson scored 35 points.",
// 球员统计数据
"A'ja Wilson 2024 season stats: 26.9 PPG, 11.9 RPG, 2.6 BPG. Won MVP award.",
"Caitlin Clark 2024 rookie stats: 19.2 PPG, 8.4 APG, 5.7 RPG. Won Rookie of the Year.",
"Breanna Stewart 2024 stats: 20.4 PPG, 8.5 RPG, 3.5 APG.",
],
[{}, {}, {}, {}, {}, {}, {}, {}, {}],
embeddings
);
const retriever = vectorStore.asRetriever({ k: 5 });
const getLatestNews = tool(
async ({ query }) => {
// 您的新闻 API 在这里
return "Latest: The WNBA announced expanded playoff format for 2025...";
},
{
name: "get_latest_news",
description: "Get the latest WNBA news and updates",
schema: z.object({ query: z.string() }),
}
);
const agent = createAgent({
model: "openai:gpt-4.1",
tools: [getLatestNews],
});
const model = new ChatOpenAI({ model: "gpt-4.1" });
const RewrittenQuery = z.object({ query: z.string() });
async function rewriteQuery(state: typeof State.State) {
const systemPrompt = `Rewrite this query to retrieve relevant WNBA information.
The knowledge base contains: team rosters, game results with scores, and player statistics (PPG, RPG, APG).
Focus on specific player names, team names, or stat categories mentioned.`;
const response = await model.withStructuredOutput(RewrittenQuery).invoke([
{ role: "system", content: systemPrompt },
{ role: "user", content: state.question },
]);
return { rewrittenQuery: response.query };
}
async function retrieve(state: typeof State.State) {
const docs = await retriever.invoke(state.rewrittenQuery);
return { documents: docs.map((doc) => doc.pageContent) };
}
async function callAgent(state: typeof State.State) {
const context = state.documents.join("\n\n");
const prompt = `Context:\n${context}\n\nQuestion: ${state.question}`;
const response = await agent.invoke({
messages: [{ role: "user", content: prompt }],
});
return { answer: response.messages.at(-1)?.contentBlocks };
}
const workflow = new StateGraph(State)
.addNode("rewrite", rewriteQuery)
.addNode("retrieve", retrieve)
.addNode("agent", callAgent)
.addEdge(START, "rewrite")
.addEdge("rewrite", "retrieve")
.addEdge("retrieve", "agent")
.addEdge("agent", END)
.compile();
const result = await workflow.invoke({
question: "Who won the 2024 WNBA Championship?",
});
console.log(result.answer);
