Update app.py

app.py

@@ -29,21 +29,22 @@ def invoke(openai_api_key, prompt, agent_option):
         raise gr.Error("Use Agent is required.")
 
     os.environ["OPENAI_API_KEY"] = openai_api_key
-    
+
+    completion = ""
     output = ""
     
     try:
         start_time_ms = round(time.time() * 1000)
         
         if (agent_option == AGENT_LANGCHAIN):
-            completion = agent_langchain(
+            result = agent_langchain(
                 config,
                 prompt
             )
     
             output = completion["output"]
         elif (agent_option == AGENT_LLAMAINDEX):
-            output = agent_llamaindex(
+            result = agent_llamaindex(
                 config,
                 prompt
             )
@@ -56,7 +57,7 @@ def invoke(openai_api_key, prompt, agent_option):
                 temperature = config["temperature"]
             )
     
-            output = completion.choices[0].message.content
+            result = completion.choices[0].message.content
     except Exception as e:
         err_msg = e
 
@@ -66,17 +67,16 @@ def invoke(openai_api_key, prompt, agent_option):
         
         trace_wandb(
             config,
-            rag_option,
+            agent_option,
             prompt, 
             completion, 
             result, 
-            callback, 
             err_msg, 
             start_time_ms, 
             end_time_ms
         )
 
-    return output
+    return result
 
 gr.close_all()