Update app.py

app.py

@@ -3,7 +3,7 @@ import os, time
 
 from dotenv import load_dotenv, find_dotenv
 
-from rag import llm_chain, rag_chain, rag_ingestion
+from rag_langchain import llm_chain, rag_chain, rag_ingestion
 from trace import trace_wandb
 
 _ = load_dotenv(find_dotenv())
@@ -16,9 +16,9 @@ config = {
     "temperature": 0            # llm
 }
 
-RAG_OFF     = "Off"
-RAG_MONGODB = "MongoDB" # serverless
-RAG_CHROMA  = "Chroma"  # requires persistent storage (small is $0.01/hour)
+RAG_OFF        = "Off"
+RAG_LANGCHAIN  = "LangChain"
+RAG_LLAMAINDEX = "LlamaIndex"
 
 def invoke(openai_api_key, prompt, rag_option):
     if (openai_api_key == ""):
@@ -42,15 +42,17 @@ def invoke(openai_api_key, prompt, rag_option):
     try:
         start_time_ms = round(time.time() * 1000)
 
-        if (rag_option == RAG_OFF):
+        if (rag_option == RAG_LANGCHAIN):
+            completion, chain, cb = rag_chain(config, rag_option, prompt)
+
+            result = completion["result"]
+        elif (rag_option == RAG_LLAMAINDEX):
+            print("TODO")
+        else:
             completion, chain, cb = llm_chain(config, prompt)
             
             if (completion.generations[0] != None and completion.generations[0][0] != None):
                 result = completion.generations[0][0].text
-        else:
-            completion, chain, cb = rag_chain(config, rag_option, prompt)
-
-            result = completion["result"]
     except Exception as e:
         err_msg = e
 
@@ -76,15 +78,15 @@ gr.close_all()
 demo = gr.Interface(fn = invoke, 
                     inputs = [gr.Textbox(label = "OpenAI API Key", type = "password", lines = 1), 
                               gr.Textbox(label = "Prompt", value = "What are GPT-4's media capabilities in 5 emojis and 1 sentence?", lines = 1),
-                              gr.Radio([RAG_OFF, RAG_MONGODB], label = "Retrieval-Augmented Generation", value = RAG_MONGODB)],
+                              gr.Radio([RAG_OFF, RAG_LANGCHAIN, RAG_LLAMAINDEX], label = "Retrieval-Augmented Generation", value = RAG_LANGCHAIN)],
                     outputs = [gr.Textbox(label = "Completion", lines = 1)],
                     title = "Context-Aware Reasoning Application",
                     description = os.environ["DESCRIPTION"],
-                    examples = [["", "What are GPT-4's media capabilities in 5 emojis and 1 sentence?", RAG_MONGODB],
-                                ["", "List GPT-4's exam scores and benchmark results.", RAG_MONGODB],
-                                ["", "Compare GPT-4 to GPT-3.5 in markdown table format.", RAG_MONGODB],
-                                ["", "Write a Python program that calls the GPT-4 API.", RAG_MONGODB],
-                                ["", "What is the GPT-4 API's cost and rate limit? Answer in English, Arabic, Chinese, Hindi, and Russian in JSON format.", RAG_MONGODB]],
+                    examples = [["", "What are GPT-4's media capabilities in 5 emojis and 1 sentence?", RAG_LANGCHAIN],
+                                ["", "List GPT-4's exam scores and benchmark results.", RAG_LANGCHAIN],
+                                ["", "Compare GPT-4 to GPT-3.5 in markdown table format.", RAG_LANGCHAIN],
+                                ["", "Write a Python program that calls the GPT-4 API.", RAG_LANGCHAIN],
+                                ["", "What is the GPT-4 API's cost and rate limit? Answer in English, Arabic, Chinese, Hindi, and Russian in JSON format.", RAG_LANGCHAIN]],
                                 cache_examples = False)
 
 demo.launch()