Merge pull request #250 from transformerlab/fix/rag-reindex

Fix RAG Reindexing and install the plugin if not already installed

Author: deep1401

transformerlab/plugins/generate_rag_outputs/main.py

@@ -137,7 +137,7 @@ async def run_evaluation():
             plugin = experiment_config.get("rag_engine")
             if plugin is None or plugin == "":
                 raise ValueError(
-                    "No RAG engine has been assigned to this experiment. Please install a RAG plugin from the Plugins Tab."
+                    "No RAG engine has been assigned to this experiment. Please install a RAG plugin from the Plugins Tab and set it by going to the Interact tab."
                 )
 
             # Set up RAG settings

transformerlab/plugins/llamaindex_simple_document_search/index.json

@@ -4,7 +4,7 @@
     "description": "",
     "plugin-format": "python",
     "type": "rag",
-    "version": "0.0.8",
+    "version": "0.0.9",
     "model_architectures": [],
     "supported_hardware_architectures": [
     "cpu",

transformerlab/plugins/llamaindex_simple_document_search/main.py

@@ -46,17 +46,26 @@ def main():
 
     args, unknown = parser.parse_known_args()
 
+    continue_after_index = False
+
     documents_dir = args.documents_dir
     persistency_dir = os.path.join(documents_dir, "persist")
 
+    if not os.path.exists(persistency_dir) and not args.index:
+        sys.stderr.write("Documents have not been indexed. Indexing them first")
+        args.index = True
+        continue_after_index = True
+
     if args.index:
         start_time = time.time()
         index_documents(documents_dir, persistency_dir, args.embedding_model)
         elapsed_time = time.time() - start_time
 
         result = {"status": "success", "elapsed_time": elapsed_time}
-        print(json.dumps(result))
-        return
+        sys.stderr.write(json.dumps(result))
+        if not continue_after_index:
+            sys.stderr.write("Indexing complete. Exiting.")
+            return
 
     # SETTINGS
     number_of_search_results = 2

transformerlab/routers/experiment/rag.py

@@ -21,9 +21,13 @@ class EmbedRequest(BaseModel):
 @router.get("/query")
 async def query(experimentId: str, query: str, settings: str = None, rag_folder: str = "rag"):
     """Query the RAG engine"""
+
     experiment_dir = await dirs.experiment_dir_by_id(experimentId)
     documents_dir = os.path.join(experiment_dir, "documents")
     documents_dir = os.path.join(documents_dir, rag_folder)
+    documents_dir = os.path.abspath(documents_dir)
+    if not documents_dir.startswith(os.path.abspath(experiment_dir)):
+        return "Error: Invalid RAG folder path"
     if not os.path.exists(documents_dir):
         return "Error: The RAG folder does not exist in the documents directory"
     experiment_details = await db.experiment_get(id=experimentId)
@@ -75,6 +79,7 @@ async def query(experimentId: str, query: str, settings: str = None, rag_folder:
         "--settings",
         settings,
     ]
+
     print(f"Calling plugin {plugin_path}" + " with model " + model + " and query " + query)
     venv_path = os.path.join(plugin_path, "venv")
     if os.path.exists(venv_path) and os.path.isdir(venv_path):
@@ -113,6 +118,7 @@ async def query(experimentId: str, query: str, settings: str = None, rag_folder:
 @router.get("/reindex")
 async def reindex(experimentId: str, rag_folder: str = "rag"):
     """Reindex the RAG engine"""
+
     experiment_dir = await dirs.experiment_dir_by_id(experimentId)
     documents_dir = os.path.join(experiment_dir, "documents")
     documents_dir = os.path.join(documents_dir, rag_folder)