refactor: add message history context | partial

Author: WillACosta

.env-example

@@ -1,4 +1,8 @@
 LANGCHAIN_API_KEY=
 LANGCHAIN_TRACING_V2=true
 LANGCHAIN_CALLBACKS_BACKGROUND=true
+
 GOOGLE_GENAI_API_KEY=
+
+REDIS_ENDPOINT_URI=localhost:6379
+REDIS_PASSWORD=redis@dev

.vscode/settings.json

@@ -1,5 +1,6 @@
 {
   "cSpell.words": [
+    "dataproviders",
     "gemini",
     "nodenext",
     "usecases"

docs/prompts/ai-prompts.md

@@ -0,0 +1,2 @@
+This document is about AI and chatbot, can you summarize it for me?
+Can you list it in a bullet list form?

package.json

@@ -23,7 +23,8 @@
     "express": "^4.21.0",
     "langchain": "^0.3.2",
     "multer": "1.4.5-lts.1",
-    "pdf-parse": "^1.1.1"
+    "pdf-parse": "^1.1.1",
+    "redis": "^4.7.0"
   },
   "devDependencies": {
     "@types/cors": "^2.8.17",

pnpm-lock.yaml

@@ -0,0 +1,11 @@
+import { createClient } from 'redis'
+
+export class RedisClient {
+	constructor() {}
+
+	_redisClient = createClient({
+		url: process.env['REDIS_ENDPOINT_URI'],
+		username: process.env['REDIS_USER'],
+		password: process.env['REDIS_PASSWORD'],
+	})
+}

src/modules/genai/adapters/dataproviders/redis/redis-client.ts

@@ -0,0 +1 @@
+export * from './dataproviders/redis/redis-client'

src/modules/genai/adapters/inde.ts

@@ -7,7 +7,10 @@ import {
 	MessagesPlaceholder,
 } from '@langchain/core/prompts'
 
-import { RunnableSequence } from '@langchain/core/runnables'
+import {
+	RunnablePassthrough,
+	RunnableSequence,
+} from '@langchain/core/runnables'
 import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter'
 import { MemoryVectorStore } from 'langchain/vectorstores/memory'
 
@@ -17,7 +20,10 @@ import {
 } from '@langchain/google-genai'
 
 import { UseCase } from '@/common/types'
-import { SEARCH_DOC_SYSTEM_PROMPT } from '../../utils'
+import {
+	CONTEXTUALIZED_SYSTEM_PROMPT,
+	SEARCH_DOC_SYSTEM_PROMPT,
+} from '../../utils'
 import { Params, Result } from './types'
 
 export class SearchInDocumentUseCase implements UseCase<Result, Params> {
@@ -30,30 +36,57 @@ export class SearchInDocumentUseCase implements UseCase<Result, Params> {
 			apiKey: process.env['GOOGLE_GENAI_API_KEY'],
 		})
 
-		const documentRetrievalChain = RunnableSequence.from([
-			(input) => input.question,
-			retriever,
-			this._convertDocsToString,
+		/// context chain init
+
+		const contextualizedPrompt = ChatPromptTemplate.fromMessages([
+			['system', CONTEXTUALIZED_SYSTEM_PROMPT],
+			new MessagesPlaceholder('chat_history'),
+			['human', '{question}'],
 		])
 
-		const answerGenerationPrompt = ChatPromptTemplate.fromMessages([
+		const contextualizedQuestionChain = RunnableSequence.from([
+			contextualizedPrompt,
+			llmModel,
+			new StringOutputParser(),
+		])
+
+		const contextualizedQuestion = (input: Record<string, any>) => {
+			if ('chat_history' in input) {
+				return contextualizedQuestionChain
+			}
+
+			return input['question']
+		}
+
+		/// context end
+
+		const qaPrompt = ChatPromptTemplate.fromMessages([
 			['system', SEARCH_DOC_SYSTEM_PROMPT],
 			new MessagesPlaceholder('chat_history'),
 			['human', '{question}'],
 		])
 
+		// if exists a history, so we need to append the history prompt
+		// to the chain
 		const retrievalChain = RunnableSequence.from([
-			{
-				context: documentRetrievalChain,
-				question: (input) => input.question,
-			},
-			answerGenerationPrompt,
+			RunnablePassthrough.assign({
+				context: (input) => {
+					if ('chat_history' in input) {
+						const chain = contextualizedQuestion(input)
+						return chain.pipe(retriever).pipe(this._convertDocsToString)
+					}
+
+					return ''
+				},
+			}),
+			qaPrompt,
 			llmModel,
 			new StringOutputParser(),
 		])
 
 		const result = await retrievalChain.invoke({
 			question: query,
+			chat_history: [], // We need to store AI messages to the history every time = result
 		})
 
 		return { result }
@@ -77,14 +110,14 @@ export class SearchInDocumentUseCase implements UseCase<Result, Params> {
 			chunkOverlap: 128,
 		})
 
-		const embeddings = new GoogleGenerativeAIEmbeddings({
-			apiKey: process.env['GOOGLE_GENAI_API_KEY'],
-			model: 'text-embedding-004',
-		})
-
 		const splitDocs = await splitter.splitDocuments(docs)
-		const vectorstore = new MemoryVectorStore(embeddings)
-		await vectorstore.addDocuments(splitDocs)
+		const vectorstore = await MemoryVectorStore.fromDocuments(
+			splitDocs,
+			new GoogleGenerativeAIEmbeddings({
+				apiKey: process.env['GOOGLE_GENAI_API_KEY'],
+				model: 'text-embedding-004',
+			}),
+		)
 
 		return vectorstore
 	}

src/modules/genai/core/usecases/search-in-document/search-in-document.usecase.ts

@@ -8,3 +8,8 @@ If you don't know the answer, just say that you don't know. Be verbose!
 
 export const REPHRASE_QUESTION_SYSTEM_TEMPLATE = `
 Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question.`
+
+export const CONTEXTUALIZED_SYSTEM_PROMPT = `Given a chat history and the latest user question
+which might reference context in the chat history, formulate a standalone question
+which can be understood without the chat history. Do NOT answer the question,
+just reformulate it if needed and otherwise return it as is.`