Update app.py

app.py

@@ -106,7 +106,7 @@ def document_retrieval_mongodb(llm, prompt):
 def llm_chain(llm, prompt):
     llm_chain = LLMChain(llm = llm, prompt = LLM_CHAIN_PROMPT)
     completion = llm_chain.run({"question": prompt})
-    return completion
+    return completion, llm_chain
 
 def rag_chain(llm, prompt, db):
     rag_chain = RetrievalQA.from_chain_type(llm, 
@@ -114,9 +114,9 @@ def rag_chain(llm, prompt, db):
                                             retriever = db.as_retriever(search_kwargs = {"k": config["k"]}), 
                                             return_source_documents = True)
     completion = rag_chain({"query": prompt})
-    return completion
+    return completion, rag_chain
 
-def wandb_trace(rag_option, prompt, completion, status_msg, start_time_ms, end_time_ms):
+def wandb_trace(rag_option, prompt, completion, chain, status_msg, start_time_ms, end_time_ms):
     wandb.init(project = "openai-llm-rag")
     if (rag_option == "Off" or str(status_msg) != ""):
         result = completion
@@ -126,7 +126,7 @@ def wandb_trace(rag_option, prompt, completion, status_msg, start_time_ms, end_t
         document_1 = completion["source_documents"][1]
         document_2 = completion["source_documents"][2]
     trace = Trace(
-        kind = "llm",
+        kind = "chain",
         name = "LLMChain" if (rag_option == "Off") else "RetrievalQA",
         status_code = "SUCCESS" if (str(status_msg) == "") else "ERROR",
         status_message = str(status_msg),
@@ -145,7 +145,8 @@ def wandb_trace(rag_option, prompt, completion, status_msg, start_time_ms, end_t
                   "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={"chain": chain}
     )
     trace.log("test")
     wandb.finish()
@@ -159,6 +160,7 @@ def invoke(openai_api_key, rag_option, prompt):
         raise gr.Error("Prompt is required.")
     completion = ""
     result = ""
+    chain = ""
     status_msg = ""
     try:
         start_time_ms = round(time.time() * 1000)
@@ -169,23 +171,23 @@ def invoke(openai_api_key, rag_option, prompt):
             #splits = document_loading_splitting()
             #document_storage_chroma(splits)
             db = document_retrieval_chroma(llm, prompt)
-            completion = rag_chain(llm, prompt, db)
+            completion, chain = rag_chain(llm, prompt, db)
             result = completion["result"]
         elif (rag_option == "MongoDB"):
             #splits = document_loading_splitting()
             #document_storage_mongodb(splits)
             db = document_retrieval_mongodb(llm, prompt)
-            completion = rag_chain(llm, prompt, db)
+            completion, chain = rag_chain(llm, prompt, db)
             result = completion["result"]
         else:
-            result = llm_chain(llm, prompt)
+            result, chain = llm_chain(llm, prompt)
             completion = result
     except Exception as e:
         status_msg = e
         raise gr.Error(e)
     finally:
         end_time_ms = round(time.time() * 1000)
-        wandb_trace(rag_option, prompt, completion, status_msg, start_time_ms, end_time_ms)
+        wandb_trace(rag_option, prompt, completion, chain, status_msg, start_time_ms, end_time_ms)
     return result
 
 gr.close_all()