Technology•January 23, 2024
Building a Wikipedia Chatbot Using Astra DB, LangChain, and Vercel
@dataclass_json @dataclass class ChunkedArticleMetadataOnly: _id: str article_metadata: ArticleMetadata chunks_metadata: dict[str, ChunkMetadata] = field(default_factory=dict) suggested_question_chunks: list[Chunk] = field(default_factory=list)
METADATA_COLLECTION.find_one_and_replace( filter={"_id": metadata._id}, replacement=metadata.to_dict(), options={"upsert": True} )
resp = METADATA_COLLECTION.find_one(filter={"_id": new_metadata._id}) prev_metadata_doc = resp["data"]["document"] prev_metadata = ChunkedArticleMetadataOnly.from_dict(prev_metadata_doc)
"use client"; import { useChat, useCompletion } from 'ai/react'; import { Message } from 'ai';
const { append, messages, isLoading, input, handleInputChange, handleSubmit } = useChat();
const handleSend = (e) => { handleSubmit(e, { options: { body: { useRag, llm, similarityMetric}}}); }
const [suggestions, setSuggestions] = useState<PromptSuggestion[]>([]); const { complete } = useCompletion({ onFinish: (prompt, completion) => { const parsed = JSON.parse(completion); const argsObj = JSON.parse(parsed?.function_call.arguments); const questions = argsObj.questions; const questionsArr: PromptSuggestion[] = []; questions.forEach(q => { questionsArr.push(q); }); setSuggestions(questionsArr); } }); useEffect(() => { complete('') }, []);
import { AstraDB } from "@datastax/astra-db-ts"; import { OpenAIStream, StreamingTextResponse } from "ai"; import OpenAI from "openai"; import type { ChatCompletionCreateParams } from 'openai/resources/chat';
const { ASTRA_DB_APPLICATION_TOKEN, ASTRA_DB_ENDPOINT, ASTRA_DB_SUGGESTIONS_COLLECTION, OPENAI_API_KEY, } = process.env; const astraDb = new AstraDB(ASTRA_DB_APPLICATION_TOKEN, ASTRA_DB_ENDPOINT); const openai = new OpenAI({ apiKey: OPENAI_API_KEY, });
const suggestionsCollection = await astraDb.collection(ASTRA_DB_SUGGESTIONS_COLLECTION); const suggestionsDoc = await suggestionsCollection.findOne( { _id: "recent_articles" }, { projection: { "recent_articles.metadata.title" : 1, "recent_articles.suggested_chunks.content" : 1, }, });
const docMap = suggestionsDoc.recent_articles.map(article => { return { pageTitle: article.metadata.title, content: article.suggested_chunks.map(chunk => chunk.content) } }); docContext = JSON.stringify(docMap);
const response = await openai.chat.completions.create({ model: "gpt-3.5-turbo-16k", stream: true, temperature: 1.5, messages: [{ role: "user", content: `You are an assistant who creates sample questions to ask a chatbot. Given the context below of the most recently added data to the most popular pages on Wikipedia come up with 4 suggested questions. Only write no more than one question per page and keep them to less than 12 words each. Do not label which page the question is for/from. START CONTEXT ${docContext} END CONTEXT `, }], functions });
const functions: ChatCompletionCreateParams.Function[] = [{ name: 'get_suggestion_and_category', description: 'Prints a suggested question and the category it belongs to.', parameters: { type: 'object', properties: { questions: { type: 'array', description: 'The suggested questions and their categories.', items: { type: 'object', properties: { category: { type: 'string', enum: ['history', 'science', 'sports', 'technology', 'arts', 'culture', 'geography', 'entertainment', 'politics', 'business', 'health'], description: 'The category of the suggested question.', }, question: { type: 'string', description: 'The suggested question.', }, }, }, }, }, required: ['questions'], }, }];
import { CohereEmbeddings } from "@langchain/cohere"; import { Document } from "@langchain/core/documents"; import { RunnableBranch, RunnableLambda, RunnableMap, RunnableSequence } from "@langchain/core/runnables"; import { StringOutputParser } from "@langchain/core/output_parsers"; import { PromptTemplate } from "langchain/prompts"; import { AstraDBVectorStore, AstraLibArgs, } from "@langchain/community/vectorstores/astradb"; import { ChatOpenAI } from "langchain/chat_models/openai"; import { StreamingTextResponse, Message } from "ai";
const questionTemplate = `You are an AI assistant answering questions about anything from Wikipedia the context will provide you with the most relevant data from wikipedia including the pages title, url, and page content. If referencing the text/context refer to it as Wikipedia. At the end of the response add one markdown link using the format: [Title](URL) and replace the title and url with the associated title and url of the more relavant page from the context This link will not be shown to the user so do not mention it. The max links you can include is 1, do not provide any other references or annotations. if the context is empty, answer it to the best of your ability. If you cannot find the answer user's question in the context, reply with "I'm sorry, I'm only allowed to answer questions related to the top 1,000 Wikipedia pages". <context> {context} </context> QUESTION: {question} `; const prompt = PromptTemplate.fromTemplate(questionTemplate);
const {messages, llm } = await req.json(); const previousMessages = messages.slice(0, -1); const latestMessage = messages[messages?.length - 1]?.content; const embeddings = new CohereEmbeddings({ apiKey: COHERE_API_KEY, inputType: "search_query", model: "embed-english-v3.0", }); const chatModel = new ChatOpenAI({ temperature: 0.5, openAIApiKey: OPENAI_API_KEY, modelName: llm ?? "gpt-4", streaming: true, });
const astraConfig: AstraLibArgs = { token: ASTRA_DB_APPLICATION_TOKEN, endpoint: ASTRA_DB_ENDPOINT, collection: “article_embeddings”, contentKey: “content” }; const vectorStore = new AstraDBVectorStore(embeddings, astraConfig); await vectorStore.initialize(); const retriever = vectorStore.asRetriever(10);
const chain = RunnableSequence.from([ condenseChatBranch, mapQuestionAndContext, prompt, chatModel, new StringOutputParser(), ]).withConfig({ runName: "chatChain"}); const stream = await chain.stream({ chat_history: formatVercelMessages(previousMessages), question: latestMessage, });
const hasChatHistoryCheck = RunnableLambda.from( (input: ChainInut) => input.chat_history.length > 0 ); const chatHistoryQuestionChain = RunnableSequence.from([ { question: (input: ChainInut) => input.question, chat_history: (input: ChainInut) => input.chat_history, }, condenseQuestionPrompt, chatModel, new StringOutputParser(), ]).withConfig({ runName: "chatHistoryQuestionChain"}); const noChatHistoryQuestionChain = RunnableLambda.from( (input: ChainInut) => input.question ).withConfig({ runName: "noChatHistoryQuestionChain"}); const condenseChatBranch = RunnableBranch.from([ [hasChatHistoryCheck, chatHistoryQuestionChain], noChatHistoryQuestionChain, ]).withConfig({ runName: "condenseChatBranch"});
const condenseQuestionTemplate = `Given the following chat history and a follow up question, If the follow up question references previous parts of the chat rephrase the follow up question to be a standalone question if not use the follow up question as the standalone question. <chat_history> {chat_history} </chat_history> Follow Up Question: {question} Standalone question:`; const condenseQuestionPrompt = PromptTemplate.fromTemplate( condenseQuestionTemplate, );
const combineDocumentsFn = (docs: Document[]) => { const serializedDocs = docs.map((doc) => `Title: ${doc.metadata.title} URL: ${doc.metadata.url} Content: ${doc.pageContent}`); return serializedDocs.join("\n\n"); }; const retrieverChain = retriever.pipe(combineDocumentsFn).withConfig({ runName: "retrieverChain"}); const mapQuestionAndContext = RunnableMap.from({ question: (input: string) => input, context: retrieverChain }).withConfig({ runName: "mapQuestionAndContext"});
return new StreamingTextResponse(stream);