Agent 设计模式——CRAG

export type GraphState = {
  /**
   * The LLM to use for the graph
   */
  llm: AzureChatOpenAI;
  /**
   * 用户查询
   */
  query: string;
  /**
   * 结果输出
   */
  result: string | null;
  /**
   * web 查询结果
   */
  needSearch: boolean;
  /**
   * 检索匹配文档
   */
  documents: Document[];
  /**
   * 阶段
   */
  steps: string[];
};

// 1. retrieval 检索节点
export async function retriever(
  state: GraphState
): Promise<Partial<GraphState>> {
  const { query } = state;
  const vectorstore = await FaissStore.load(directory, embeddingModel);
  const retriever = vectorstore.asRetriever(8);
 
  // rag 上下文检索
  const contextRetrieverChain = RunnableSequence.from([
    // 传递参数
    (input) => input.query,
    // 检索返回文档
    retriever,
  ]);
 
  const result = await contextRetrieverChain.invoke({
    query,
  });
 
  return {
    documents: result,
    steps: ["retriever"],
  };
}
 
// 2. 对检索结果进行评估，打分
export async function gradeDocuments(
  state: GraphState
): Promise<Partial<GraphState>> {
  const { llm, query, documents, steps } = state;
 
  const gradePrompt = ChatPromptTemplate.fromMessages([
    [
      "system",
      `You are a teacher grading a quiz. You will be given: 
1/ a QUESTION 
2/ a set of comma separated FACTS provided by the student
 
You are grading RELEVANCE RECALL:
A score of 1 means that ANY of the FACTS are relevant to the QUESTION. 
A score of 0 means that NONE of the FACTS are relevant to the QUESTION. 
1 is the highest (best) score. 0 is the lowest score you can give. 
 
Explain your reasoning in a step-by-step manner. Ensure your reasoning and conclusion are correct. 
 
Avoid simply stating the correct answer at the outset. \n\n {format_instructions}`,
    ],
    ["user", `FACTS: \n\n {document} \n\n QUESTION: {query}`],
  ]);
 
  const gradeSchema = z.object({
    score: z
      .string()
      .describe(`Binary score for Documents relevant to question`),
  });
  const parser = StructuredOutputParser.fromZodSchema(gradeSchema);
 
  const gradeChain = gradePrompt.pipe(llm).pipe(parser);
 
  // 挑选匹配文档，非匹配文档调用查询工具
  const filterDocuments = [];
  let needSearch = false;
 
  for (const doc of documents) {
    const result = await gradeChain.invoke({
      query,
      document: doc.pageContent,
      format_instructions: parser.getFormatInstructions(),
    });
    const { score } = result;
    if (score === "1") {
      filterDocuments.push(doc);
    } else {
      needSearch = true;
      continue;
    }
  }
 
  return {
    steps: steps.concat("gradeDocuments"),
    documents: filterDocuments,
    needSearch,
  };
}
 
// 3. 调用工具，搜索公开信息
const tavilySearchTool = new TavilySearchResults();
 
export async function webSearch(
  state: GraphState
): Promise<Partial<GraphState>> {
  const { query, documents, steps } = state;
 
  const result = await tavilySearchTool.invoke(query);
 
  let searchResult = [];
  try {
    const searchJSON = JSON.parse(result);
    searchResult = searchJSON.map((d: any) => ({
      pageContent: d.content,
      metadata: { url: d.url },
    }));
  } catch (err) {
    console.error(`parse err: `, err);
    searchResult = [];
  }
 
  return {
    steps: steps.concat("webSearch"),
    documents: documents.concat(searchResult),
  };
}
 
const prompt = ChatPromptTemplate.fromMessages([
  [
    "system",
    `
    You are an assistant for question-answering tasks. 
    Use the following documents to answer the question.
    If you don't know the answer, just say that you don't know.
  `,
  ],
  [
    "user",
    `
    Documents: {documents} 
    Question: {query}
    Answer:
    `,
  ],
]);
 
// 4. 完成 RAG，输出结果
export async function generation(
  state: GraphState
): Promise<Partial<GraphState>> {
  const { llm, query, documents } = state;
 
  // 构建 Rag 链条
  const ragChain = RunnableSequence.from([
    // 将 context 和 query 传递给 prompt
    prompt,
    // 构建完毕的 prompt 传递给 llm
    llm,
    // 解析生成文本
    new StringOutputParser(),
  ]);
 
  const result = await ragChain.invoke({
    query,
    documents: documents.map((d) => d.pageContent).join(`\n\n`),
  });
 
  return {
    result,
    steps: state.steps.concat(["generation"]),
  };
}

const validateSearch = (state: GraphState): "webSearch" | "generation" => {
  return state.needSearch === true ? "webSearch" : "generation";
};
 
function createGraph() {
  const graph = new StateGraph<GraphState>({
    channels: graphChannels,
  });
 
  graph.addNode("retriever", retriever);
  graph.addNode("gradeDocuments", gradeDocuments);
  graph.addNode("webSearch", webSearch);
  graph.addNode("generation", generation);
 
  graph.addEdge(START, "retriever");
  graph.addEdge("retriever", "gradeDocuments");
  graph.addEdge("webSearch", "generation");
  graph.addConditionalEdges("gradeDocuments", validateSearch);
  graph.addEdge("generation", END);
 
  const app = graph.compile();
  return app;
}
 
async function main(query: string) {
  const app = createGraph();
 
  await app.stream({
    llm,
    query,
  });
}
 
// 测试一：标准场景
main("What are the types of agent memory?");
 
// 测试二：异常 case
// main("Who won the 2024 NBA championShip?");