Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-import logging, os, sys, time
+import logging, os, sys, threading, time
 
 from dotenv import load_dotenv, find_dotenv
 
@@ -7,6 +7,8 @@ from rag_langchain import LangChainRAG
 from rag_llamaindex import LlamaIndexRAG
 from trace import trace_wandb
 
+lock = threading.Lock()
+
 _ = load_dotenv(find_dotenv())
 
 RAG_INGESTION = False # load, split, embed, and store documents
@@ -34,55 +36,58 @@ def invoke(openai_api_key, prompt, rag_option):
     if (rag_option is None):
         raise gr.Error("Retrieval-Augmented Generation is required.")
 
-    os.environ["OPENAI_API_KEY"] = openai_api_key
-    
-    if (RAG_INGESTION):
-        if (rag_option == RAG_LANGCHAIN):
-            rag = LangChainRAG()
-            rag.ingestion(config)
-        elif (rag_option == RAG_LLAMAINDEX):
-            rag = LlamaIndexRAG()
-            rag.ingestion(config)
-
-    completion = ""
-    result = ""
-    callback = ""
-    err_msg = ""
+    with lock:
+        os.environ["OPENAI_API_KEY"] = openai_api_key
+        
+        if (RAG_INGESTION):
+            if (rag_option == RAG_LANGCHAIN):
+                rag = LangChainRAG()
+                rag.ingestion(config)
+            elif (rag_option == RAG_LLAMAINDEX):
+                rag = LlamaIndexRAG()
+                rag.ingestion(config)
     
-    try:
-        start_time_ms = round(time.time() * 1000)
-
-        if (rag_option == RAG_LANGCHAIN):
-            rag = LangChainRAG()
-            completion, callback = rag.rag_chain(config, prompt)
-            result = completion["result"]
-        elif (rag_option == RAG_LLAMAINDEX):
-            rag = LlamaIndexRAG()
-            result, callback = rag.retrieval(config, prompt)
-        else:
-            rag = LangChainRAG()
-            completion, callback = rag.llm_chain(config, prompt)
-            result = completion.generations[0][0].text
-    except Exception as e:
-        err_msg = e
-
-        raise gr.Error(e)
-    finally:
-        end_time_ms = round(time.time() * 1000)
+        completion = ""
+        result = ""
+        callback = ""
+        err_msg = ""
         
-        trace_wandb(
-            config,
-            rag_option,
-            prompt, 
-            completion, 
-            result, 
-            callback, 
-            err_msg, 
-            start_time_ms, 
-            end_time_ms
-        )
-
-    return result
+        try:
+            start_time_ms = round(time.time() * 1000)
+    
+            if (rag_option == RAG_LANGCHAIN):
+                rag = LangChainRAG()
+                completion, callback = rag.rag_chain(config, prompt)
+                result = completion["result"]
+            elif (rag_option == RAG_LLAMAINDEX):
+                rag = LlamaIndexRAG()
+                result, callback = rag.retrieval(config, prompt)
+            else:
+                rag = LangChainRAG()
+                completion, callback = rag.llm_chain(config, prompt)
+                result = completion.generations[0][0].text
+        except Exception as e:
+            err_msg = e
+    
+            raise gr.Error(e)
+        finally:
+            end_time_ms = round(time.time() * 1000)
+            
+            trace_wandb(
+                config,
+                rag_option,
+                prompt, 
+                completion, 
+                result, 
+                callback, 
+                err_msg, 
+                start_time_ms, 
+                end_time_ms
+            )
+
+            del os.environ["OPENAI_API_KEY"]
+    
+        return result
 
 gr.close_all()