Update app.py

app.py

@@ -123,24 +123,25 @@ def wandb_trace(rag_option, prompt, completion, result, chain, err_msg, start_ti
     wandb.init(project = "openai-llm-rag")
     trace = Trace(
         kind = "chain",
-        name = type(chain).__name__ if (chain != None) else "",
+        name = "" if (chain == None) else type(chain).__name__,
         status_code = "success" if (str(err_msg) == "") else "error",
         status_message = str(err_msg),
         metadata = {
-            "chunk_overlap": "" if (rag_option == RAG_OFF) else config["chunk_overlap"],
-            "chunk_size": "" if (rag_option == RAG_OFF) else config["chunk_size"],
-            #"k": "" if (rag_option == RAG_OFF) else config["k"],
-            #"model_name": config["model_name"],
-            #"temperature": config["temperature"],
+            "chunk_overlap": config["chunk_overlap"] if (str(err_msg) == "" and rag_option != RAG_OFF) else "",
+            "chunk_size": config["chunk_size"] if (str(err_msg) == "" and rag_option != RAG_OFF) else "",
         },
         inputs = {"rag_option": rag_option if (str(err_msg) == "") else "",
-                  "prompt": str(prompt if (str(err_msg) == "") else ""),
+                  "prompt": prompt if (str(err_msg) == "") else "",
         },
         outputs = {"result": result if (str(err_msg) == "") else "",
-                   "completion": str(completion),
+                   "completion": str(completion) if (str(err_msg) == "") else "",
         },
-        model_dict = {"llm_client": str(chain.llm.client) if (rag_option == RAG_OFF) else str(chain.combine_documents_chain.llm_chain.llm.client),
-                      "chain_prompt": str(chain.prompt) if (rag_option == RAG_OFF) else str(chain.combine_documents_chain.llm_chain.prompt),
+        model_dict = {"llm_client": str(chain.llm.client) if (rag_option == RAG_OFF) else 
+                                    str(chain.combine_documents_chain.llm_chain.llm.client),
+                      "model_name": config["model_name"],
+                      "temperature": config["temperature"],
+                      "chain_prompt": str(chain.prompt) if (rag_option == RAG_OFF) else 
+                                      str(chain.combine_documents_chain.llm_chain.prompt),
                       "retriever": "" if (rag_option == RAG_OFF) else str(chain.retriever),
         },
         start_time_ms = start_time_ms,
@@ -177,26 +178,10 @@ def invoke(openai_api_key, rag_option, prompt):
             db = document_retrieval_mongodb(llm, prompt)
             completion, chain = rag_chain(llm, prompt, db)
             result = completion["result"]
-            #print("###")
-            #print(completion)
-            #print("###")
-            #print(chain.combine_documents_chain)
-            print("###")
-            print(chain.combine_documents_chain.llm_chain.llm.client)
-            print("###")
-            print(chain.combine_documents_chain.llm_chain.prompt)
-            print("###")
-            print(chain.retriever)
-            print("###")
         else:
             completion, chain = llm_chain(llm, prompt)
             result = completion.generations[0][0].text if (completion.generations[0] != None and 
                                                            completion.generations[0][0] != None) else ""
-            print("###")
-            print(chain.llm.client)
-            print("###")
-            print(chain.prompt)
-            print("###")
     except Exception as e:
         err_msg = e
         raise gr.Error(e)