feat: Update frontend to use new RAG endpoint (#113)

Author: yoomlam

app/src/app_config.py

@@ -11,6 +11,11 @@
 
 class AppConfig(PydanticBaseEnvConfig):
     environment: str = "local"
+    # Use the similar syntax as the preview environment bucket names
+    # Note: since there's a single ChromaDB instance for all environments,
+    # there can be collision if multiple developers are running RAG simultaneously,
+    # but this is unlikely at this time.
+    bucket_name: str = "local-labs-referral-pilot-app-dev"
     # Set HOST to 127.0.0.1 by default to avoid other machines on the network
     # from accessing the application. This is especially important if you are
     # running the application locally on a public network. This needs to be
@@ -63,8 +68,9 @@ def chroma_client(self) -> ClientAPI:
         return chromadb.HttpClient(host=self.rag_db_host, port=self.rag_db_port)
 
     def chroma_document_store(self) -> ChromaDocumentStore:
+        # Use bucket name as part of collection name to avoid collision with `dev` environment for PRs
         return ChromaDocumentStore(
-            collection_name=f"{self.collection_name_prefix}_{self.environment}",
+            collection_name=f"{self.collection_name_prefix}_{self.bucket_name}",
             host=self.rag_db_host,
             port=self.rag_db_port,
         )

app/src/ingestion/rag_utils.py

@@ -3,11 +3,13 @@
 from pathlib import Path
 
 from botocore.exceptions import NoCredentialsError
+from chromadb.api import ClientAPI
 from haystack import Pipeline
 from haystack.components.converters import MultiFileConverter
 from haystack.components.embedders import SentenceTransformersDocumentEmbedder
 from haystack.components.preprocessors import DocumentPreprocessor
 from haystack.components.writers import DocumentWriter
+from haystack.document_stores.errors.errors import DocumentStoreError
 from haystack_integrations.document_stores.chroma import ChromaDocumentStore
 
 from src.app_config import config
@@ -16,19 +18,40 @@
 logger = logging.getLogger(__name__)
 
 
+def delete_preview_collections(chroma_client: ClientAPI) -> None:
+    collections = chroma_client.list_collections()
+    for collection in collections:
+        name = collection.name
+        if name.startswith(f"{config.collection_name_prefix}_p-"):
+            logger.info("Deleting preview collection: %s", name)
+            chroma_client.delete_collection(name)
+
+
 def populate_vector_db() -> None:
     logging.basicConfig(format="%(levelname)s - %(name)s -  %(message)s", level=logging.INFO)
 
     chroma_client = config.chroma_client()
+    delete_preview_collections(chroma_client)
+
     logger.info("ChromaDB collections: %s", chroma_client.list_collections())
     doc_store = config.chroma_document_store()
     collection_name = doc_store._collection_name
 
     # Clear existing collection if any
     if doc_store.count_documents() > 0:
-        # Don't delete collection since it's referenced by existing pipelines upon their startup
-        logger.info("Clearing out existing vector DB collection=%r", collection_name)
-        doc_store.delete_all_documents()
+        try:
+            # Don't delete collection since it's referenced by existing pipelines upon their startup
+            logger.info("Clearing out existing vector DB collection=%r", collection_name)
+            # recreate_index=True results in a new id for the collection, which breaks existing pipelines
+            doc_store.delete_all_documents(recreate_index=False)
+        except DocumentStoreError as e:
+            # Ignore this error from haystack.logging, which is okay since logging is the last step in delete_all_documents
+            if "overwrite 'name' in LogRecord" in str(e):
+                logger.info("Ignoring expected DocumentStoreError: %s", e)
+            else:
+                logger.warning("Unexpected DocumentStoreError: %s", e)
+                raise
+    assert doc_store.count_documents() == 0, "Documents should be deleted from collection"
 
     # Download files from S3
     local_folder = download_s3_folder_to_local()

frontend/src/util/fetchResources.ts

@@ -7,7 +7,13 @@ export async function fetchResources(
   prompt_version_id: string | null,
 ) {
   const apiDomain = await getApiDomain();
-  const url = apiDomain + "generate_referrals/run";
+
+  const useNonRag = process.env.NEXT_PUBLIC_USE_NONRAG === "true";
+  const url_path = useNonRag
+    ? "generate_referrals/run"
+    : "generate_referrals_rag/run";
+
+  const url = apiDomain + url_path;
   const headers = {
     "Content-Type": "application/json",
   };