Update app.py

app.py

@@ -104,7 +104,7 @@ def document_retrieval_mongodb(llm, prompt):
     return db
 
 def llm_chain(llm, prompt):
-    llm_chain = LLMChain(llm = llm, prompt = LLM_CHAIN_PROMPT)
+    llm_chain = LLMChain(llm = llm, prompt = LLM_CHAIN_PROMPT, verbose = True)
     completion = llm_chain.run({"question": prompt})
     return completion, llm_chain
 
@@ -112,11 +112,15 @@ def rag_chain(llm, prompt, db):
     rag_chain = RetrievalQA.from_chain_type(llm, 
                                             chain_type_kwargs = {"prompt": RAG_CHAIN_PROMPT}, 
                                             retriever = db.as_retriever(search_kwargs = {"k": config["k"]}), 
-                                            return_source_documents = True)
+                                            return_source_documents = True,
+                                            verbose = True)
     completion = rag_chain({"query": prompt})
     return completion, rag_chain
 
 def wandb_trace(rag_option, prompt, completion, chain, status_msg, start_time_ms, end_time_ms):
+    #print(chain.inputKey)
+    #print(chain.outputKey)
+    #print(chain.retriever)
     wandb.init(project = "openai-llm-rag")
     if (rag_option == "Off" or str(status_msg) != ""):
         result = completion
@@ -145,8 +149,7 @@ def wandb_trace(rag_option, prompt, completion, chain, status_msg, start_time_ms
                   "document_2": "" if (rag_option == "Off" or str(status_msg) != "") else str(document_2)},
         outputs = {"result": result},
         start_time_ms = start_time_ms,
-        end_time_ms = end_time_ms,
-        model_dict={"chain": chain}
+        end_time_ms = end_time_ms
     )
     trace.log("test")
     wandb.finish()