Update app.py

app.py

@@ -69,26 +69,31 @@ def document_storage_chroma(splits):
                                       persist_directory = CHROMA_DIR)
 
 def document_storage_mongodb(splits):
+    #TODO
     vector_db = Chroma.from_documents(documents = splits, 
                                       embedding = OpenAIEmbeddings(disallowed_special = ()), 
                                       persist_directory = CHROMA_DIR)
 
 def document_retrieval_chroma(llm, prompt):
-    vector_db = Chroma(embedding_function = OpenAIEmbeddings(),
+    db = Chroma(embedding_function = OpenAIEmbeddings(),
                        persist_directory = CHROMA_DIR)
-    rag_chain = RetrievalQA.from_chain_type(llm, 
-                                            chain_type_kwargs = {"prompt": RAG_CHAIN_PROMPT}, 
-                                            retriever = vector_db.as_retriever(search_kwargs = {"k": 3}), 
-                                            return_source_documents = True)
-    result = rag_chain({"query": prompt})
-    return result["result"]
+    return db
 
 def document_retrieval_mongodb(llm, prompt):
-    vector_db = Chroma(embedding_function = OpenAIEmbeddings(),
-                       persist_directory = CHROMA_DIR)
+    #TODO    
+    db = Chroma(embedding_function = OpenAIEmbeddings(),
+                persist_directory = CHROMA_DIR)
+    return db
+
+def llm_chain(llm, prompt):
+    llm_chain = LLMChain(llm = llm, prompt = LLM_CHAIN_PROMPT)
+    result = llm_chain.run({"question": prompt})
+    return result
+
+def rag_chain(llm, prompt, db):
     rag_chain = RetrievalQA.from_chain_type(llm, 
                                             chain_type_kwargs = {"prompt": RAG_CHAIN_PROMPT}, 
-                                            retriever = vector_db.as_retriever(search_kwargs = {"k": 3}), 
+                                            retriever = db.as_retriever(search_kwargs = {"k": 3}), 
                                             return_source_documents = True)
     result = rag_chain({"query": prompt})
     return result["result"]
@@ -107,13 +112,14 @@ def invoke(openai_api_key, rag_option, prompt):
         #splits = document_loading_splitting()
         if (rag_option == "Chroma"):
             #document_storage_chroma(splits)
-            result = document_retrieval_chroma(llm, prompt)
+            db = document_retrieval_chroma(llm, prompt)
+            result = rag_chain(llm, prompt, db)
         elif (rag_option == "MongoDB"):
             #document_storage_mongodb(splits)
-            result = document_retrieval_mongodb(llm, prompt)
+            db = document_retrieval_mongodb(llm, prompt)
+            result = rag_chain(llm, prompt, db)
         else:
-            chain = LLMChain(llm = llm, prompt = LLM_CHAIN_PROMPT)
-            result = chain.run({"question": prompt})
+            result = llm_chain(llm, prompt)
     except Exception as e:
         raise gr.Error(e)
     return result