refactor: improve gen-ai module architecture and separate logic]

Author: WillACosta

src/di/di_modules.ts

@@ -0,0 +1,30 @@
+import {
+  ChatMemory,
+  DocumentsService,
+  LLMService,
+} from '@/modules/genai/adapters'
+
+import {
+  SearchInDocumentUseCase,
+  TranslateTextUseCase,
+} from '@/modules/genai/core'
+
+const chatMemory = new ChatMemory()
+const llmService = new LLMService()
+const documentService = new DocumentsService(llmService)
+
+const searchInDocumentUseCase = new SearchInDocumentUseCase(
+	chatMemory,
+	llmService,
+	documentService,
+)
+
+const translateUseCase = new TranslateTextUseCase()
+
+export {
+  chatMemory,
+  documentService,
+  llmService,
+  searchInDocumentUseCase,
+  translateUseCase
+}

src/modules/genai/adapters/dataproviders/chat-memory/chat-memory.ts

@@ -0,0 +1,34 @@
+import { RedisChatMessageHistory } from '@langchain/redis'
+import { BufferMemory } from 'langchain/memory'
+
+export class ChatMemory {
+	constructor() {
+		this._memory = new BufferMemory({
+			chatHistory: new RedisChatMessageHistory({
+				sessionId: 'a168c61a-c431-4ef8-bc1c-fedd808d45ea',
+				sessionTTL: 3600,
+				config: {
+					url: process.env['REDIS_URL'],
+					password: process.env['REDIS_PASSWORD'],
+				},
+			}),
+			returnMessages: true,
+			memoryKey: 'chat_history',
+		})
+	}
+
+	private _memory: BufferMemory
+
+	get memoryObject() {
+		return this._memory
+	}
+
+	async retrieveMemoryHistory(): Promise<Record<string, any>[]> {
+		const memoryResults = await this._memory.loadMemoryVariables({})
+		return memoryResults['chat_history']
+	}
+
+	saveChatHistory(input: string, output: string): Promise<void> {
+		return this._memory.saveContext({ input }, { output })
+	}
+}

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

@@ -21,9 +21,9 @@ export class RedisClient {
 		await this._redisClient.quit()
 	}
 
-	async storeValue(key: string, value: object): Promise<void> {
+	async storeValue(key: string, value: object, ttl?: number): Promise<void> {
 		await this._redisClient.set(key, JSON.stringify(value), {
-			EX: 300,
+			EX: ttl,
 			NX: true,
 		})
 	}

src/modules/genai/adapters/inde.ts

@@ -0,0 +1,5 @@
+export * from './dataproviders/chat-memory/chat-memory'
+export * from './dataproviders/redis/redis-client'
+export * from './services/document/documents.service'
+export * from './services/llm/llm.service'
+

src/modules/genai/adapters/index.ts

@@ -0,0 +1,40 @@
+import { PDFLoader } from '@langchain/community/document_loaders/fs/pdf'
+import { Document } from 'langchain/document'
+import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter'
+import { MemoryVectorStore } from 'langchain/vectorstores/memory'
+import { LLMService } from '../llm/llm.service'
+
+export class DocumentsService {
+	constructor(private _llmService: LLMService) {}
+
+	convertDocsToString(documents: Document[]) {
+		return documents
+			.map((document) => `<doc>\n${document.pageContent}\n</doc>`)
+			.join('\n')
+	}
+
+	async initializeVectorStore(docs: Document<Record<string, any>>[]) {
+		const splitter = new RecursiveCharacterTextSplitter({
+			chunkSize: 1536,
+			chunkOverlap: 128,
+		})
+
+		const splitDocs = await splitter.splitDocuments(docs)
+		const vectorstore = await MemoryVectorStore.fromDocuments(
+			splitDocs,
+			this._llmService.textEmbedding,
+		)
+
+		const retriever = vectorstore.asRetriever()
+		return { vectorstore, retriever }
+	}
+
+	async loadDocument(filePath: string) {
+		const systemPath = process.cwd()
+		const loader = new PDFLoader(`${systemPath}${filePath}`, {
+			parsedItemSeparator: '',
+		})
+
+		return await loader.load()
+	}
+}

src/modules/genai/adapters/services/document/documents.service.ts

@@ -0,0 +1,22 @@
+import {
+	ChatGoogleGenerativeAI,
+	GoogleGenerativeAIEmbeddings,
+} from '@langchain/google-genai'
+
+export class LLMService {
+	constructor() {}
+
+	get llm() {
+		return new ChatGoogleGenerativeAI({
+			model: 'gemini-1.5-flash',
+			apiKey: process.env['GOOGLE_GENAI_API_KEY'],
+		})
+	}
+
+	get textEmbedding() {
+		return new GoogleGenerativeAIEmbeddings({
+			model: 'text-embedding-004',
+			apiKey: process.env['GOOGLE_GENAI_API_KEY'],
+		})
+	}
+}

src/modules/genai/adapters/services/llm/llm.service.ts

@@ -1,18 +1,13 @@
-import { Request, Response } from 'express'
+import { Request, Response } from 'express';
 
-import {
-	SearchInDocumentUseCase,
-	TranslateTextUseCase,
-	upload,
-} from '../../core'
+import { searchInDocumentUseCase, translateUseCase } from '@/di/di_modules';
+import { upload } from '../../core';
 
 export class GenAIController {
 	async translateText(
 		req: Request<any, any, { text: string; language: string }>,
 		res: Response,
 	) {
-		// TODO: Implement DI
-		const translateUseCase = new TranslateTextUseCase()
 		const { text, language } = req.body
 
 		if (language === undefined || text === undefined) {
@@ -29,8 +24,6 @@ export class GenAIController {
 		req: Request<any, any, { query: string }>,
 		res: Response,
 	) {
-		const useCase = new SearchInDocumentUseCase()
-
 		upload(req, res, async (err) => {
 			if (err) {
 				return res
@@ -45,7 +38,10 @@ export class GenAIController {
 			const query = req.body.query
 			const filePath = `/uploads/${req.file.filename}`
 
-			const { result } = await useCase.invoke({ query, filePath })
+			const { result } = await searchInDocumentUseCase.invoke({
+				query,
+				filePath,
+			})
 
 			return res.json({
 				message: 'File and query received successfully!',

src/modules/genai/application/controllers/gen-ai.controller.ts

@@ -1,5 +1,3 @@
-import { PDFLoader } from '@langchain/community/document_loaders/fs/pdf'
-import { Document } from '@langchain/core/documents'
 import { StringOutputParser } from '@langchain/core/output_parsers'
 
 import {
@@ -11,50 +9,34 @@ import {
 	RunnablePassthrough,
 	RunnableSequence,
 } from '@langchain/core/runnables'
-import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter'
-import { MemoryVectorStore } from 'langchain/vectorstores/memory'
-
-import {
-	ChatGoogleGenerativeAI,
-	GoogleGenerativeAIEmbeddings,
-} from '@langchain/google-genai'
-
-import { UseCase } from '@/common/types'
-import { RedisChatMessageHistory } from '@langchain/redis'
-import { BufferMemory } from 'langchain/memory'
 
 import {
 	CONTEXTUALIZED_SYSTEM_PROMPT,
 	SEARCH_DOC_SYSTEM_PROMPT,
 } from '../../utils'
 
+import { UseCase } from '@/common/types'
 import { Params, Result } from './types'
 
-/// TODO: delegate to RedisClient class
-const chatMemory = new BufferMemory({
-	chatHistory: new RedisChatMessageHistory({
-		sessionId: 'a168c61a-c431-4ef8-bc1c-fedd808d45ea',
-		sessionTTL: 3600, // 300 = 5m, 3600 = 1h, null = never
-		config: {
-			url: process.env['REDIS_URL'],
-			password: process.env['REDIS_PASSWORD'],
-		},
-	}),
-	returnMessages: true,
-	memoryKey: 'chat_history',
-})
+import {
+	ChatMemory,
+	DocumentsService,
+	LLMService,
+} from '@/modules/genai/adapters'
 
 export class SearchInDocumentUseCase implements UseCase<Result, Params> {
-	async invoke({ filePath, query }: Params): Promise<Result> {
-		const docs = await this._loadDocument(filePath)
-		const vectorStore = await this._initializeVectorStoreWithDocuments(docs)
-		const retriever = vectorStore.asRetriever()
-		const llmModel = new ChatGoogleGenerativeAI({
-			model: 'gemini-1.5-flash',
-			apiKey: process.env['GOOGLE_GENAI_API_KEY'],
-		})
+	constructor(
+		private _memory: ChatMemory,
+		private _llmService: LLMService,
+		private _documentService: DocumentsService,
+	) {}
 
-		/// Standalone question referencing past context
+	async invoke({ filePath, query }: Params): Promise<Result> {
+		const llmModel = this._llmService.llm
+		const docs = await this._documentService.loadDocument(filePath)
+		const { retriever } = await this._documentService.initializeVectorStore(
+			docs,
+		)
 
 		const contextualizedPrompt = ChatPromptTemplate.fromMessages([
 			['system', CONTEXTUALIZED_SYSTEM_PROMPT],
@@ -68,8 +50,6 @@ export class SearchInDocumentUseCase implements UseCase<Result, Params> {
 			new StringOutputParser(),
 		])
 
-		/// Standalone question end
-
 		const questionAnsweringPrompt = ChatPromptTemplate.fromMessages([
 			['system', SEARCH_DOC_SYSTEM_PROMPT],
 			new MessagesPlaceholder('chat_history'),
@@ -83,7 +63,7 @@ export class SearchInDocumentUseCase implements UseCase<Result, Params> {
 						return RunnableSequence.from([
 							contextualizedQuestionChain,
 							retriever,
-							this._convertDocsToString,
+							this._documentService.convertDocsToString,
 						])
 					}
 
@@ -95,50 +75,13 @@ export class SearchInDocumentUseCase implements UseCase<Result, Params> {
 			new StringOutputParser(),
 		])
 
-		const memoryResults = await chatMemory.loadMemoryVariables({})
-		const history = memoryResults['chat_history']
+		const history = await this._memory.retrieveMemoryHistory()
 		const result = await retrievalChain.invoke({
 			question: query,
 			chat_history: history,
 		})
 
-		await chatMemory.saveContext({ input: query }, { output: result })
+		this._memory.saveChatHistory(query, result)
 		return { result }
 	}
-
-	private async _loadDocument(filePath: string) {
-		const systemPath = process.cwd()
-		const loader = new PDFLoader(`${systemPath}${filePath}`, {
-			parsedItemSeparator: '',
-		})
-
-		return await loader.load()
-	}
-
-	// TODO: improve it and delegate vector store to another layer
-	private async _initializeVectorStoreWithDocuments(
-		docs: Document<Record<string, any>>[],
-	) {
-		const splitter = new RecursiveCharacterTextSplitter({
-			chunkSize: 1536,
-			chunkOverlap: 128,
-		})
-
-		const splitDocs = await splitter.splitDocuments(docs)
-		const vectorstore = await MemoryVectorStore.fromDocuments(
-			splitDocs,
-			new GoogleGenerativeAIEmbeddings({
-				apiKey: process.env['GOOGLE_GENAI_API_KEY'],
-				model: 'text-embedding-004',
-			}),
-		)
-
-		return vectorstore
-	}
-
-	private _convertDocsToString(documents: Document[]) {
-		return documents
-			.map((document) => `<doc>\n${document.pageContent}\n</doc>`)
-			.join('\n')
-	}
 }

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

@@ -1,9 +1,17 @@
-export const SEARCH_DOC_SYSTEM_PROMPT = `You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Be verbose!
+export const SEARCH_DOC_SYSTEM_PROMPT = `You are an assistant for question-answering tasks.
+Rules:
+
+- Use the following pieces of retrieved context to answer the question.
+- If you don't know the answer, just say that you don't know.
+- If the use ask something about a past conversation, you can check the chat history and analyze the content for answer the user properly.
 
 <context>
 {context}
 </context>`
 
 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.`
+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.`