Update app.py

app.py

@@ -116,31 +116,14 @@ def rag_chain(llm, prompt, db):
     return completion, rag_chain
 
 def wandb_trace(rag_option, prompt, completion, chain, status_msg, start_time_ms, end_time_ms):
-    if (chain != None):
-        if (type(chain).__name__ == "LLMChain"):
-            print("1=" + str(chain.llm.client))
-            print("1=" + str(chain.llm.async_client))
-            print("1=" + str(chain.llm.model_name))
-            print("1=" + str(chain.llm.temperature))
-            print("2=" + str(chain.prompt.input_variables))
-            print("2=" + str(chain.prompt.template))
-        #if (type(chain).__name__ == "RetrievalQA"):
-            #print(chain.combine_documents_chain.llm_chain.llm)
-            #print(chain.combine_documents_chain.llm_chain.async_client)
-            #print(chain.combine_documents_chain.retriever)
-            #print(chain.combine_documents_chain.vectorstore)
-            #print("3=" + str(chain.llm))
-            #print("4=" + str(chain.chain_type_kwargs))
-            #print("5=" + str(chain.retriever))
     wandb.init(project = "openai-llm-rag")
     if (rag_option == "Off" or str(status_msg) != ""):
         result = completion
     else:
         result = completion["result"]
-        documents = completion["source_documents"]
-        document_0 = completion["source_documents"][0].metadata["source"]
-        document_1 = completion["source_documents"][1].metadata["source"]
-        document_2 = completion["source_documents"][2].metadata["source"]
+        doc_meta_source_0 = completion["source_documents"][0].metadata["source"]
+        doc_meta_source_1 = completion["source_documents"][1].metadata["source"]
+        doc_meta_source_2 = completion["source_documents"][2].metadata["source"]
     trace = Trace(
         kind = "chain",
         name = type(chain).__name__ if (chain != None) else "",
@@ -156,14 +139,12 @@ def wandb_trace(rag_option, prompt, completion, chain, status_msg, start_time_ms
         inputs = {"rag_option": rag_option if (str(status_msg) == "") else "",
                   "prompt": str(prompt if (str(status_msg) == "") else ""), 
                   "prompt_template": str((llm_template if (rag_option == "Off") else rag_template) if (str(status_msg) == "") else ""),
-                  "documents": "" if (rag_option == "Off" or str(status_msg) != "") else str(documents),
-                  "document_0": "" if (rag_option == "Off" or str(status_msg) != "") else str(document_0),
-                  "document_1": "" if (rag_option == "Off" or str(status_msg) != "") else str(document_1),
-                  "document_2": "" if (rag_option == "Off" or str(status_msg) != "") else str(document_2)},
+                  "doc_meta_source_0": "" if (rag_option == "Off" or str(status_msg) != "") else str(doc_meta_source_0),
+                  "doc_meta_source_1": "" if (rag_option == "Off" or str(status_msg) != "") else str(doc_meta_source_1),
+                  "doc_meta_source_2": "" if (rag_option == "Off" or str(status_msg) != "") else str(doc_meta_source_2},
         outputs = {"result": result},
         start_time_ms = start_time_ms,
-        end_time_ms = end_time_ms#,
-        #model_dict = {"": x, "": x}
+        end_time_ms = end_time_ms
     )
     trace.log("test")
     wandb.finish()