Update app.py

app.py

@@ -3,7 +3,8 @@ import os, time
 
 from dotenv import load_dotenv, find_dotenv
 
-from rag_langchain import llm_chain, rag_chain, rag_ingestion
+from rag_langchain import llm_chain, rag_chain, rag_ingestion_langchain
+from rag_llamaindex import rag_ingestion_llamaindex, rag_retrieval
 from trace import trace_wandb
 
 _ = load_dotenv(find_dotenv())
@@ -31,25 +32,28 @@ def invoke(openai_api_key, prompt, rag_option):
     os.environ["OPENAI_API_KEY"] = openai_api_key
     
     if (RAG_INGESTION):
-        rag_ingestion(config)
+        if (rag_option == RAG_LANGCHAIN):
+            rag_ingestion_llangchain(config)
+        elif (rag_option == RAG_LLAMAINDEX):
+            rag_ingestion_llamaindex(config)
     
     chain = None
     completion = ""
     result = ""
-    cb = ""
+    callback = ""
     err_msg = ""
     
     try:
         start_time_ms = round(time.time() * 1000)
 
         if (rag_option == RAG_LANGCHAIN):
-            completion, chain, cb = rag_chain(config, prompt)
+            completion, chain, callback = rag_chain(config, prompt)
 
             result = completion["result"]
         elif (rag_option == RAG_LLAMAINDEX):
-            print("TODO")
+            result = rag_retrieval(config, prompt)
         else:
-            completion, chain, cb = llm_chain(config, prompt)
+            completion, chain, callback = llm_chain(config, prompt)
             
             if (completion.generations[0] != None and completion.generations[0][0] != None):
                 result = completion.generations[0][0].text
@@ -66,7 +70,7 @@ def invoke(openai_api_key, prompt, rag_option):
         #            completion, 
         #            result, 
         #            chain, 
-        #            cb, 
+        #            callback, 
         #            err_msg, 
         #            start_time_ms, 
         #            end_time_ms)