Update rag_utils.py

rag_utils.py

@@ -1,27 +1,20 @@
-# rag_utils.py
-
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import FAISS
 from langchain.chains import RetrievalQA
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_groq import ChatGroq
+from langchain.docstore.document import Document
 
-def create_vectorstore_from_text(text: str):
-    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-    texts = splitter.split_text(text)
-    
-    embeddings = HuggingFaceEmbeddings(
-        model_name="sentence-transformers/all-MiniLM-L6-v2",
-        model_kwargs={"device": "cpu"}
-    )
+def create_vectorstore_from_text(documents, embeddings):
+    # If string is passed instead of list of Document, convert it
+    if isinstance(documents, str):
+        splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+        chunks = splitter.split_text(documents)
+        documents = [Document(page_content=chunk) for chunk in chunks]
     
-    vectorstore = FAISS.from_texts(texts, embedding=embeddings)
+    vectorstore = FAISS.from_documents(documents, embedding=embeddings)
     return vectorstore
 
-def create_rag_chain(vectorstore):
+def create_rag_chain(llm, vectorstore):
     retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
-    
-    llm = ChatGroq(model_name="llama3-8b-8192", temperature=0)
-    
-    rag_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever)
-    return rag_chain
+    return RetrievalQA.from_chain_type(llm=llm, retriever=retriever)