Add Faiss db RAG

Signed-off-by: Jagadish Krishnamoorthy <[email protected]>

app.py

@@ -1,17 +1,34 @@
+from collections import defaultdict
+
 import streamlit as st
-from langchain_openai import ChatOpenAI
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_community.document_loaders import WebBaseLoader
+from langchain_text_splitters import CharacterTextSplitter
+from langchain_community.vectorstores import FAISS
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain.chains import create_retrieval_chain
 
 default_url = "https://rocm.docs.amd.com/en/latest/what-is-rocm.html"
 st.title("URL Loader")
 
+embeddings = OpenAIEmbeddings()
+
 url = st.text_input("Provide URL ", default_url)
 if "url_dict" not in st.session_state:
     st.session_state.url_dict = {}
 if url not in st.session_state.url_dict:
     loader = WebBaseLoader(url)
-    st.session_state.url_dict[url] = loader.load()
+    documents = loader.load()
+    st.session_state.url_dict[url] = defaultdict(dict)
+    st.session_state.url_dict[url]['documents'] =  documents
+    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
+    docs = text_splitter.split_documents(documents)
+    db = FAISS.from_documents(docs, embeddings)
+    print(db.index.ntotal)
+    url_hash = "faiss_index" + str(abs(hash(url)))
+    db.save_local(url_hash)
+    st.session_state.url_dict[url]['FAISS_db'] = url_hash
 
 llm = ChatOpenAI(temperature=0.1)
 prompt = ChatPromptTemplate.from_template("""
@@ -19,31 +36,37 @@ prompt = ChatPromptTemplate.from_template("""
     Context: {context}
     Question : {input}
     """)
-chain = prompt | llm
 
 # Initialize chat history
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
 # Display chat messages from history on app rerun
-for message in st.session_state.messages:
+for message in st.session_state.messages[-2:]:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
 
 # React to user input
-if prompt := st.chat_input("Ask Question to the URL provided"):
+if question := st.chat_input("Ask Question to the URL provided"):
     # Display user message in chat message container
-    st.chat_message("user").markdown(prompt)
+    st.chat_message("user").markdown(question)
     # Add user message to chat history
-    st.session_state.messages.append({"role": "user", "content": prompt})
+    st.session_state.messages.append({"role": "user", "content": question})
+
+    db = FAISS.load_local(st.session_state.url_dict[url]['FAISS_db'],
+                              embeddings, allow_dangerous_deserialization=True)
+    document_chain = create_stuff_documents_chain(
+        llm=llm,
+        prompt=prompt
+    )
+    retriever = db.as_retriever(search_kwargs={"k": 2})
+    chain = create_retrieval_chain(retriever, document_chain)
 
     response = chain.invoke({
-    "context" : [st.session_state.url_dict[url]],
-    "input" : prompt
+    "input" : question
     })
     # Display assistant response in chat message container
     with st.chat_message("assistant"):
-        st.markdown(response.content)
-    # Add assistant response to chat history
-    st.session_state.messages.append({"role": "assistant", "content": response.content})
+        st.markdown(response["answer"])
+    st.session_state.messages.append({"role": "assistant", "content": response["answer"]})
 

requirements.txt

@@ -13,6 +13,7 @@ charset-normalizer==3.3.2
 click==8.1.7
 dataclasses-json==0.6.6
 distro==1.9.0
+faiss-cpu==1.8.0
 frozenlist==1.4.1
 gitdb==4.0.11
 GitPython==3.1.43