Update app.py

app.py

@@ -117,7 +117,13 @@ def rag_chain(llm, prompt, db):
 
 def wandb_trace(rag_option, prompt, completion, chain, status_msg, start_time_ms, end_time_ms):
     if (chain != None):
-        print(type(chain).__name__)
+        if (type(chain).__name__ == "LLMChain")
+        print("1=" + chain.llm)
+        print("2=" + chain.prompt)
+        if (type(chain).__name__ == "RetrievalQA")
+        print("3=" + chain.llm)
+        print("4=" + chain.chain_type_kwargs)
+        print("5=" + chain.retriever)
     wandb.init(project = "openai-llm-rag")
     if (rag_option == "Off" or str(status_msg) != ""):
         result = completion
@@ -128,7 +134,7 @@ def wandb_trace(rag_option, prompt, completion, chain, status_msg, start_time_ms
         document_2 = completion["source_documents"][2]
     trace = Trace(
         kind = "chain",
-        name = "LLMChain" if (rag_option == "Off") else "RetrievalQA",
+        name = type(chain).__name__ if (chain != None) else "",
         status_code = "SUCCESS" if (str(status_msg) == "") else "ERROR",
         status_message = str(status_msg),
         metadata={
@@ -146,7 +152,8 @@ 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
+        end_time_ms = end_time_ms,
+        model_dict = {"": x, "": x}
     )
     trace.log("test")
     wandb.finish()