Update custom_llm.py

custom_llm.py

@@ -7,13 +7,30 @@ import requests
 from langchain.prompts import PromptTemplate, ChatPromptTemplate
 from operator import itemgetter
 
-# from langchain.memory import ConversationBufferMemory
 from langchain.memory import ChatMessageHistory, ConversationBufferMemory
 from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_community.chat_models import ChatOpenAI
 from langchain_core.runnables import RunnableLambda, RunnablePassthrough
 from langchain_core.messages import AIMessage, HumanMessage
 
+from langchain_community.document_loaders import DirectoryLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import PyMuPDFLoader
+import os
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+
+
+def create_vectorstore():
+    loader = os.getenv('knowledge_base')
+    
+    splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=20)
+    
+    docs = splitter.create_documents([loader])
+    
+    emb_model = HuggingFaceEmbeddings(model_name='sentence-transformers/paraphrase-multilingual-mpnet-base-v2', encode_kwargs={'normalize_embeddings': True})
+    db = FAISS.from_documents(docs, emb_model)
+    return db
 
 def custom_chain_with_history(llm, memory):
 
@@ -50,7 +67,7 @@ def custom_chain_with_history(llm, memory):
     #     ]
     # )
     
-    return {"chat_history":prompt_memory, "context":db.as_retriever(search_type="similarity", search_kwargs={"k": 8}) | format_docs, "question": RunnablePassthrough()} | prompt | llm
+    return {"chat_history":prompt_memory, "context":create_vectorstore().as_retriever(search_type="similarity", search_kwargs={"k": 8}) | format_docs, "question": RunnablePassthrough()} | prompt | llm
 
 class CustomLLM(LLM):
     repo_id : str