Upload rag_util.py

rag_util.py

@@ -0,0 +1,66 @@
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import PyPDFLoader
+# from langchain_core.vectorstores import InMemoryVectorStore
+# from langchain_community.vectorstores import Chroma
+from langchain_chroma import Chroma
+from langchain_huggingface import HuggingFaceEmbeddings
+
+
+embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+# vector_database = InMemoryVectorStore(embedding_model)
+
+# vector_database = Chroma(
+#     persist_directory="./chroma_store",
+#     embedding_function=embedding_model
+# )
+
+
+
+
+
+# document_retriever = vector_database.as_retriever(search_type="mmr", search_kwargs={"k" : 3, "lambda_mult": 0.8})
+
+
+def load_pdf_document(file_path):
+    document_loader = PyPDFLoader(file_path)
+    return document_loader.load()
+
+def chunk_documents(raw_documents):
+    text_processor = RecursiveCharacterTextSplitter(
+        chunk_size = 1000,
+        chunk_overlap = 200,
+        add_start_index = True
+    )
+    return text_processor.split_documents(raw_documents)
+
+def find_related_documents(query, vector_database):
+    # return vector_database.similarity_search(query, k=2)
+    return vector_database.max_marginal_relevance_search(query, k=2, fetch_k=5, lambda_mult=0.6)
+
+
+def ProcessDocuments(document_path: str, chatID: str) -> str:
+
+    loaded_doc = load_pdf_document(document_path)
+    chunked_doc = chunk_documents(loaded_doc)
+
+    
+    vector_database = Chroma(
+        persist_directory=f"./chroma_store/{chatID}",
+        embedding_function=embedding_model
+    )
+
+    vector_database.add_documents(chunked_doc)
+
+
+def Create_RAG_Prompt(query: str, chatID: str):
+    
+    vector_database = Chroma(
+        persist_directory=f"./chroma_store/{chatID}",
+        embedding_function=embedding_model
+    )
+
+    relevant_docs = find_related_documents(query, vector_database)
+    context_text = "\n\n".join([doc.page_content for doc in relevant_docs])
+ 
+    return query, context_text
+