Modular code

app.py

@@ -27,14 +27,18 @@ chain = prompt | gemini
 
 index_name = "langchain-test-index"
 
-def store_embeddings(pdf_path, chunk_size, chunk_overlap):
+def extract_text_from_pdf(pdf_path):
     raw_documents = []
     for path in pdf_path:
         raw_documents.extend(PyPDFLoader(path).load())
+    return raw_documents
 
+def chunk_text(raw_documents):
     text_splitter = CharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
     documents = text_splitter.split_documents(raw_documents)
+    return documents
 
+def delete_pinecone():
     pc = Pinecone(api_key=os.environ["PINECONE_API_KEY"])
 
     index = pc.Index(host="https://langchain-test-index-la2n80y.svc.aped-4627-b74a.pinecone.io")
@@ -42,13 +46,53 @@ def store_embeddings(pdf_path, chunk_size, chunk_overlap):
     if index.describe_index_stats()['total_vector_count'] > 0:
         index.delete(delete_all=True)
 
+def store_chroma_db(documents):
     chroma_db = Chroma.from_documents(documents, embeddings, persist_directory="./chroma_db")
+
+def store_faiss_db(documents):
     faiss_db = FAISS.from_documents(documents, embeddings)
     faiss_db.save_local("./faiss_db")
+
+def store_lance_db(documents):
     lance_db = LanceDB.from_documents(documents, embeddings, uri="./lance_db")
+
+def store_pinecone_db(documents):
     pinecone_db = PineconeVectorStore.from_documents(documents, index_name=index_name,
                                                      embedding=embeddings)
     
+def load_chroma_db():
+    chroma_db = Chroma(persist_directory="./chroma_db", embedding_function=embeddings)
+    return chroma_db
+
+def load_faiss_db():
+    faiss_db = FAISS.load_local("./faiss_db", embeddings, allow_dangerous_deserialization=True)
+    return faiss_db
+
+def load_lance_db():
+    lance_db = LanceDB(embedding=embeddings, uri="./lance_db")
+    return lance_db
+
+def connect_pinecone_db():
+    pinecone_db = PineconeVectorStore(index_name=index_name, embedding=embeddings)
+    return pinecone_db
+
+def invoke_chain(db, query):
+    docs = db.similarity_search(query)
+    answer = chain.invoke({"context":docs, "question": query}, return_only_outputs=True)
+    return answer
+    
+def store_embeddings(pdf_path, chunk_size, chunk_overlap):
+    raw_documents = extract_text_from_pdf(pdf_path)
+
+    documents = chunk_text(raw_documents)
+
+    delete_pinecone()
+
+    store_chroma_db(documents)
+    store_chroma_db(documents)
+    store_lance_db(documents)
+    store_pinecone_db(documents)
+    
     return "All embeddings are stored in vector database"
 
 title = "PDF Chat"
@@ -57,21 +101,17 @@ examples = [[["data/amazon-10-k-2024.pdf"], 1000, 100],
             [["data/goog-10-k-2023.pdf"], 1000, 100]]
 
 def inference(query):
-    chroma_db = Chroma(persist_directory="./chroma_db", embedding_function=embeddings)
-    chroma_docs = chroma_db.similarity_search(query)
-    chroma_answer = chain.invoke({"context":chroma_docs, "question": query}, return_only_outputs=True)
+    chroma_db = load_chroma_db()
+    chroma_answer = invoke_chain(chroma_db, query)
 
-    faiss_db = FAISS.load_local("./faiss_db", embeddings, allow_dangerous_deserialization=True)
-    faiss_docs = faiss_db.similarity_search(query)
-    faiss_answer = chain.invoke({"context":faiss_docs, "question": query}, return_only_outputs=True)
+    faiss_db = load_faiss_db()
+    faiss_answer = invoke_chain(faiss_db, query)
 
-    lance_db = LanceDB(embedding=embeddings, uri="./lance_db")
-    lance_docs = lance_db.similarity_search(query)
-    lance_answer = chain.invoke({"context":lance_docs, "question": query}, return_only_outputs=True)
+    lance_db = load_lance_db()
+    lance_answer = invoke_chain(lance_db, query)
 
-    pinecone_db = PineconeVectorStore(index_name=index_name, embedding=embeddings)
-    pinecone_docs = pinecone_db.similarity_search(query)
-    pinecoce_answer = chain.invoke({"context":pinecone_docs, "question": query}, return_only_outputs=True)
+    pinecone_db = connect_pinecone_db()
+    pinecoce_answer = invoke_chain(pinecone_db, query)
 
     return chroma_answer, faiss_answer, lance_answer, pinecoce_answer