Update app.py

app.py

@@ -3,7 +3,14 @@
 import gradio as gr
 import logging, os, sys, threading
 
-from custom_utils import connect_to_database, rag_ingestion, rag_retrieval, rag_inference
+from custom_utils import (
+    connect_to_database,
+    inference,
+    rag_ingestion,
+    rag_retrieval_naive,
+    rag_retrieval_advanced,
+    rag_inference
+)
 
 lock = threading.Lock()
 
@@ -31,22 +38,42 @@ def invoke(openai_api_key,
     with lock:
         db, collection = connect_to_database()
 
+        inference_result = ""
+        
         if (RAG_INGESTION):
             return rag_ingestion(collection)
-        else:
-            retrieval_result = rag_retrieval(openai_api_key, 
-                                             prompt, 
-                                             accomodates,
-                                             bedrooms,
-                                             db, 
-                                             collection)
-            inference_result = rag_inference(openai_api_key, 
-                                             prompt, 
-                                             retrieval_result)
-            print("###")
-            print(inference_result)
-            print("###")
-            return inference_result
+        elif rag_option == RAGOFF:
+            inference_result = inference(
+                openai_api_key, 
+                prompt)
+        elif rag_option == RAG_NAIVE:
+            retrieval_result = rag_retrieval_naive(
+                openai_api_key, 
+                prompt,
+                db, 
+                collection)
+            inference_result = rag_inference(
+                openai_api_key, 
+                prompt, 
+                retrieval_result)        
+        elif rag_option == RAG_ADVANCED:
+            retrieval_result = rag_retrieval_advanced(
+                openai_api_key, 
+                prompt, 
+                accomodates,
+                bedrooms,
+                db, 
+                collection)
+            inference_result = rag_inference(
+                openai_api_key, 
+                prompt, 
+                retrieval_result)
+
+        #print("###")
+        #print(inference_result)
+        #print("###")
+ 
+        return inference_result
 
 gr.close_all()
 
@@ -56,7 +83,7 @@ demo = gr.Interface(
               gr.Textbox(label = "Prompt", value = os.environ["PROMPT"], lines = 1),
               gr.Number(label = "Accomodates", value = 2),
               gr.Number(label = "Bedrooms", value = 1),
-              gr.Radio([RAG_OFF, RAG_NAIVE, RAG_ADVANCED], label = "Retrieval-Augmented Generation", value = RAG_NAIVE)],
+              gr.Radio([RAG_OFF, RAG_NAIVE, RAG_ADVANCED], label = "Retrieval-Augmented Generation", value = RAG_ADVANCED)],
     outputs = [gr.Markdown(label = "Completion", value = os.environ["COMPLETION"], line_breaks = True, sanitize_html = False)],
     title = "Context-Aware Reasoning Application",
     description = os.environ["DESCRIPTION"]