Update app.py

app.py

@@ -119,8 +119,21 @@ def rag_chain(llm, prompt, db):
     completion = rag_chain({"query": prompt})
     return completion, rag_chain
 
-def wandb_trace(rag_option, prompt, completion, result, generation_info, llm_output, chain, err_msg, start_time_ms, end_time_ms):
+def wandb_trace(rag_option, prompt, completion, chain, err_msg, start_time_ms, end_time_ms):
+    result = ""
+    generation_info = ""
+    llm_output = ""
+    
+    if (rag_option == RAG_OFF):
+        if (completion.generations[0] != None and completion.generations[0][0] != None):
+            result = completion.generations[0][0].text
+            generation_info = completion.generations[0][0].generation_info
+        llm_output = completion.llm_output
+    else:
+        result = completion["result"]
+    
     wandb.init(project = "openai-llm-rag")
+    
     trace = Trace(
         kind = "chain",
         name = "" if (chain == None) else type(chain).__name__,
@@ -130,9 +143,9 @@ def wandb_trace(rag_option, prompt, completion, result, generation_info, llm_out
             "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": prompt if (str(err_msg) == "") else "",
-        },
+        inputs = {"rag_option": rag_option,
+                  "prompt": prompt,
+        } if (str(err_msg) == "") else "",
         outputs = {"result": result if (str(err_msg) == "") else "",
                    "generation_info": str(generation_info) if (str(err_msg) == "") else "",
                    "llm_output": str(llm_output) if (str(err_msg) == "") else "",
@@ -151,6 +164,7 @@ def wandb_trace(rag_option, prompt, completion, result, generation_info, llm_out
         start_time_ms = start_time_ms,
         end_time_ms = end_time_ms
     )
+    
     trace.log("test")
     wandb.finish()
 
@@ -161,41 +175,35 @@ def invoke(openai_api_key, rag_option, prompt):
         raise gr.Error("Retrieval Augmented Generation is required.")
     if (prompt == ""):
         raise gr.Error("Prompt is required.")
+    
     chain = None
     completion = ""
-    result = ""
-    generation_info = ""
-    llm_output = ""
     err_msg = ""
+    
     try:
         start_time_ms = round(time.time() * 1000)
         llm = ChatOpenAI(model_name = config["model_name"], 
                          openai_api_key = openai_api_key, 
                          temperature = config["temperature"])
+        
         if (rag_option == RAG_CHROMA):
             #splits = document_loading_splitting()
             #document_storage_chroma(splits)
             db = document_retrieval_chroma(llm, prompt)
             completion, chain = rag_chain(llm, prompt, db)
-            result = completion["result"]
         elif (rag_option == RAG_MONGODB):
             #splits = document_loading_splitting()
             #document_storage_mongodb(splits)
             db = document_retrieval_mongodb(llm, prompt)
             completion, chain = rag_chain(llm, prompt, db)
-            result = completion["result"]
         else:
             completion, chain = llm_chain(llm, prompt)
-            if (completion.generations[0] != None and completion.generations[0][0] != None):
-                result = completion.generations[0][0].text
-                generation_info = completion.generations[0][0].generation_info
-            llm_output = completion.llm_output
     except Exception as e:
         err_msg = e
         raise gr.Error(e)
     finally:
         end_time_ms = round(time.time() * 1000)
-        wandb_trace(rag_option, prompt, completion, result, generation_info, llm_output, chain, err_msg, start_time_ms, end_time_ms)
+        wandb_trace(rag_option, prompt, completion, chain, err_msg, start_time_ms, end_time_ms)
     return result
 
 gr.close_all()