import {
AzureDocumentDBVectorStore,
AzureDocumentDBSimilarityType,
} from "@langchain/azure-cosmosdb";
import { ChatPromptTemplate } from "@langchain/core/prompts";
import { ChatOpenAI, OpenAIEmbeddings } from "@langchain/openai";
import { createStuffDocumentsChain } from "@langchain/classic/chains/combine_documents";
import { createRetrievalChain } from "@langchain/classic/chains/retrieval";
import { TextLoader } from "@langchain/classic/document_loaders/fs/text";
import { RecursiveCharacterTextSplitter } from "@langchain/textsplitters";
// 从文件加载文档
const loader = new TextLoader("./state_of_the_union.txt");
const rawDocuments = await loader.load();
const splitter = new RecursiveCharacterTextSplitter({
chunkSize: 1000,
chunkOverlap: 0,
});
const documents = await splitter.splitDocuments(rawDocuments);
// 创建 Azure DocumentDB 向量存储
const store = await AzureDocumentDBVectorStore.fromDocuments(
documents,
new OpenAIEmbeddings(),
{
databaseName: "langchain",
collectionName: "documents",
indexOptions: {
numLists: 100,
dimensions: 1536,
similarity: AzureDocumentDBSimilarityType.COS,
},
}
);
// 执行相似性搜索
const resultDocuments = await store.similaritySearch(
"总统对凯坦吉·布朗·杰克逊说了什么？"
);
console.log("相似性搜索结果：");
console.log(resultDocuments[0].pageContent);
/*
今晚。我呼吁参议院：通过《自由投票法案》。通过《约翰·刘易斯投票权法案》。并且，在你们处理这些的同时，通过《披露法案》，以便美国人能够知道谁在资助我们的选举。
今晚，我想向一位毕生致力于服务这个国家的人致敬：斯蒂芬·布雷耶大法官——一位陆军退伍军人、宪法学者，以及即将退休的美国最高法院大法官。布雷耶大法官，感谢您的服务。
总统最严肃的宪法职责之一是提名某人在美国最高法院任职。
我在四天前就这样做了，当时我提名了巡回上诉法院法官凯坦吉·布朗·杰克逊。她是我们国家顶尖的法律头脑之一，将继续布雷耶大法官卓越的遗产。
*/
// 将存储用作链的一部分
const model = new ChatOpenAI({ model: "gpt-3.5-turbo-1106" });
const questionAnsweringPrompt = ChatPromptTemplate.fromMessages([
[
"system",
"根据以下上下文回答用户的问题：\n\n{context}",
],
["human", "{input}"],
]);
const combineDocsChain = await createStuffDocumentsChain({
llm: model,
prompt: questionAnsweringPrompt,
});
const chain = await createRetrievalChain({
retriever: store.asRetriever(),
combineDocsChain,
});
const res = await chain.invoke({
input: "总统在价格方面的首要任务是什么？",
});
console.log("链响应：");
console.log(res.answer);
/*
总统的首要任务是控制价格。
*/
// 清理
await store.delete();
await store.close();
