Update app.py

app.py

@@ -119,19 +119,7 @@ def rag_chain(llm, prompt, db):
     completion = rag_chain({"query": prompt})
     return completion, rag_chain
 
-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"]
-    
+def wandb_trace(rag_option, prompt, completion, result, generation_info, llm_output, chain, err_msg, start_time_ms, end_time_ms):
     wandb.init(project = "openai-llm-rag")
     
     trace = Trace(
@@ -178,6 +166,9 @@ def invoke(openai_api_key, rag_option, prompt):
     
     chain = None
     completion = ""
+    result = ""
+    generation_info = ""
+    llm_output = ""
     err_msg = ""
     
     try:
@@ -191,19 +182,25 @@ def invoke(openai_api_key, rag_option, prompt):
             #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, chain, err_msg, start_time_ms, end_time_ms)
+        wandb_trace(rag_option, prompt, completion, result, generation_info, llm_output, chain, err_msg, start_time_ms, end_time_ms)
     return result
 
 gr.close_all()