OpenSearch support (#223)

* add http auth for the low level client as well

* wip on chunk retrieval

* implement raw query to return the node contents for OpenSearch

* don't refetch chat history on window focus to allow for continued streaming

* add in oss username and password

* drop databases

lastFile:llm-service/app/services/amp_metadata/__init__.py

* create opensearch config for ui and backend

* wip

lastFile:ui/src/pages/Settings/VectorDBFields.tsx

* drop databases

lastFile:ui/src/pages/Settings/VectorDBFields.tsx

* WIP vector db provider

lastFile:ui/src/pages/Settings/VectorDBFields.tsx

* cleanup opensearch config setting

* add namespace for opensearch

* Handle 404s properly with opensearch deletions & other possibly failed deletion

* fix a small mypy issue

* add beta tag

* Update ui/src/pages/RagChatTab/SessionsSidebar/CreateSession/CreateSessionForm.tsx

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* handle the case where we can't find the chunk

* use typeadapters

* log query

---------

Co-authored-by: Elijah Williams <ewilliams@cloudera.com>
Co-authored-by: Michael Liu <mliu@cloudera.com>
Co-authored-by: Baasit Sharief <baasitsharief@gmail.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 5 people

.env.example

@@ -13,6 +13,12 @@ OPENAI_API_VERSION=
 # QDRANT or OPENSEARCH
 VECTOR_DB_PROVIDER=QDRANT
 
+# OpenSearch
+OPENSEARCH_ENDPOINT=
+OPENSEARCH_USERNAME=
+OPENSEARCH_PASSWORD=
+OPENSEARCH_NAMESPACE=
+
 # AWS
 AWS_ACCESS_KEY_ID=
 AWS_SECRET_ACCESS_KEY=

backend/src/main/java/com/cloudera/cai/rag/external/RagBackendClient.java

@@ -45,6 +45,7 @@
 import com.cloudera.cai.util.Tracker;
 import com.cloudera.cai.util.exceptions.ClientError;
 import com.cloudera.cai.util.exceptions.HttpError;
+import com.cloudera.cai.util.exceptions.NotFound;
 import com.cloudera.cai.util.exceptions.ServerError;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.core.JsonProcessingException;
@@ -53,9 +54,11 @@
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.List;
+import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Component;
 
+@Slf4j
 @Component
 public class RagBackendClient {
   private static final String AUTH_TOKEN = System.getenv("CDSW_APIV2_KEY");
@@ -125,22 +128,34 @@ public String createSummary(Types.RagDocument ragDocument, String bucketName) {
   }
 
   public void deleteDataSource(Long dataSourceId) {
-    client.delete(
-        getLlmServiceUrl() + "/data_sources/" + dataSourceId,
-        "Authorization",
-        "Bearer " + AUTH_TOKEN);
+    try {
+      client.delete(
+          getLlmServiceUrl() + "/data_sources/" + dataSourceId,
+          "Authorization",
+          "Bearer " + AUTH_TOKEN);
+    } catch (NotFound e) {
+      log.info("Data source not found. Deletion not necessary.");
+    }
   }
 
   public void deleteDocument(long dataSourceId, String documentId) {
-    client.delete(
-        getLlmServiceUrl() + "/data_sources/" + dataSourceId + "/documents/" + documentId,
-        "Authorization",
-        "Bearer " + AUTH_TOKEN);
+    try {
+      client.delete(
+          getLlmServiceUrl() + "/data_sources/" + dataSourceId + "/documents/" + documentId,
+          "Authorization",
+          "Bearer " + AUTH_TOKEN);
+    } catch (NotFound e) {
+      log.info("Document not found. Deletion not necessary.");
+    }
   }
 
   public void deleteSession(Long sessionId) {
-    client.delete(
-        getLlmServiceUrl() + "/sessions/" + sessionId, "Authorization", "Bearer " + AUTH_TOKEN);
+    try {
+      client.delete(
+          getLlmServiceUrl() + "/sessions/" + sessionId, "Authorization", "Bearer " + AUTH_TOKEN);
+    } catch (NotFound e) {
+      log.info("Session not found. Deletion not necessary.");
+    }
   }
 
   record IndexRequest(

backend/src/main/java/com/cloudera/cai/util/SimpleHttpClient.java

@@ -109,6 +109,10 @@ public void delete(String path, String... headers) {
       HttpResponse<String> response =
           httpClient.send(request, HttpResponse.BodyHandlers.ofString());
       int statusCode = response.statusCode();
+      if (statusCode == 404) {
+        throw new NotFound("Failed to delete. Not Found");
+      }
+
       if (statusCode >= 400) {
         throw new RuntimeException(
             "Failed to delete " + path + " code: " + statusCode + ", body : " + response.body());

backend/src/test/java/com/cloudera/cai/rag/external/RagBackendClientTest.java

@@ -213,6 +213,20 @@ void deleteDataSource() {
                 HttpMethod.DELETE, "http://localhost:8081/data_sources/1234", null));
   }
 
+  @Test
+  void deleteDataSource_notFound() {
+    Tracker<TrackedHttpRequest<?>> tracker = new Tracker<>();
+    RagBackendClient client =
+        new RagBackendClient(SimpleHttpClient.createNull(tracker, new NotFound("Not found!")));
+    client.deleteDataSource(1234L);
+    List<TrackedHttpRequest<?>> values = tracker.getValues();
+    assertThat(values)
+        .hasSize(1)
+        .contains(
+            new TrackedHttpRequest<>(
+                HttpMethod.DELETE, "http://localhost:8081/data_sources/1234", null));
+  }
+
   @Test
   void deleteDocument() {
     Tracker<TrackedHttpRequest<?>> tracker = new Tracker<>();
@@ -228,6 +242,22 @@ void deleteDocument() {
                 null));
   }
 
+  @Test
+  void deleteDocument_notFound() {
+    Tracker<TrackedHttpRequest<?>> tracker = new Tracker<>();
+    RagBackendClient client =
+        new RagBackendClient(SimpleHttpClient.createNull(tracker, new NotFound("Not found!")));
+    client.deleteDocument(1234L, "documentId");
+    List<TrackedHttpRequest<?>> values = tracker.getValues();
+    assertThat(values)
+        .hasSize(1)
+        .contains(
+            new TrackedHttpRequest<>(
+                HttpMethod.DELETE,
+                "http://localhost:8081/data_sources/1234/documents/documentId",
+                null));
+  }
+
   @Test
   void deleteSession() {
     Tracker<TrackedHttpRequest<?>> tracker = new Tracker<>();
@@ -241,6 +271,20 @@ void deleteSession() {
                 HttpMethod.DELETE, "http://localhost:8081/sessions/1234", null));
   }
 
+  @Test
+  void deleteSession_notFound() {
+    Tracker<TrackedHttpRequest<?>> tracker = new Tracker<>();
+    RagBackendClient client =
+        new RagBackendClient(SimpleHttpClient.createNull(tracker, new NotFound("Not found!")));
+    client.deleteSession(1234L);
+    List<TrackedHttpRequest<?>> values = tracker.getValues();
+    assertThat(values)
+        .hasSize(1)
+        .contains(
+            new TrackedHttpRequest<>(
+                HttpMethod.DELETE, "http://localhost:8081/sessions/1234", null));
+  }
+
   @Test
   void null_handlesThrowable() {
     RagBackendClient client =

llm-service/app/ai/vector_stores/opensearch.py

@@ -40,7 +40,10 @@
 from abc import ABC
 from typing import Optional, List
 
+import fastapi.exceptions
+import opensearchpy
 from llama_index.core.base.embeddings.base import BaseEmbedding
+from llama_index.core.schema import BaseNode, TextNode
 from llama_index.core.vector_stores.types import BasePydanticVectorStore
 from llama_index.vector_stores.opensearch import (
     OpensearchVectorStore,
@@ -55,18 +58,21 @@
 
 logger = logging.getLogger(__name__)
 
+
 def _new_opensearch_client(dim: int, index: str) -> OpensearchVectorClient:
     return OpensearchVectorClient(
-        # username=os.environ.get("OPENSEARCH_USERNAME", "admin"),
-        # password=os.environ.get("OPENSEARCH_INITIAL_ADMIN_PASSWORD"),
         endpoint=settings.opensearch_endpoint,
         index=index,
         dim=dim,
+        http_auth=(settings.opensearch_username, settings.opensearch_password),
     )
 
 
 def _get_low_level_client() -> OpensearchClient:
-    os_client = OpensearchClient(settings.opensearch_endpoint)
+    os_client = OpensearchClient(
+        settings.opensearch_endpoint,
+        http_auth=(settings.opensearch_username, settings.opensearch_password),
+    )
     return os_client
 
 
@@ -77,14 +83,14 @@ class OpenSearch(VectorStore, ABC):
     def for_chunks(data_source_id: int) -> "OpenSearch":
         return OpenSearch(
             data_source_id=data_source_id,
-            table_name=f"index_{data_source_id}",
+            table_name=f"{settings.opensearch_namespace}__index_{data_source_id}",
         )
 
     @staticmethod
     def for_summaries(data_source_id: int) -> "OpenSearch":
         return OpenSearch(
             data_source_id=data_source_id,
-            table_name=f"summary_index_{data_source_id}",
+            table_name=f"{settings.opensearch_namespace}__summary_index_{data_source_id}",
         )
 
     def __init__(
@@ -110,7 +116,10 @@ def size(self) -> Optional[int]:
 
     def delete(self) -> None:
         os_client = self._low_level_client
-        os_client.indices.delete(index=self.table_name)
+        try:
+            os_client.indices.delete(index=self.table_name)
+        except opensearchpy.exceptions.NotFoundError:
+            raise fastapi.exceptions.HTTPException(404, "Index not found")
 
     def delete_document(self, document_id: str) -> None:
         self._get_client().delete_by_doc_id(document_id)
@@ -120,6 +129,19 @@ def llama_vector_store(self) -> BasePydanticVectorStore:
             self._get_client(),
         )
 
+    def get_chunk_contents(self, chunk_id: str) -> BaseNode:
+        query = {"query": {"terms": {"_id": [chunk_id]}}}
+        raw_results = self._low_level_client.search(index=self.table_name, body=query)
+        if raw_results["hits"] and raw_results["hits"]["hits"]:
+            return TextNode(
+                id_=chunk_id, text=raw_results["hits"]["hits"][0]["_source"]["content"]
+            )
+        else:
+            logger.error(f"Chunk not found for query: {query}")
+            raise fastapi.exceptions.HTTPException(
+                404, "Chunk not found with id: " + chunk_id
+            )
+
     def _get_client(self) -> OpensearchVectorClient:
         return _new_opensearch_client(
             dim=self._find_dim(self.data_source_id),

llm-service/app/ai/vector_stores/vector_store.py

@@ -41,6 +41,7 @@
 
 import umap
 from llama_index.core.base.embeddings.base import BaseEmbedding
+from llama_index.core.schema import BaseNode
 from llama_index.core.vector_stores.types import BasePydanticVectorStore
 
 logger = logging.getLogger(__name__)
@@ -109,3 +110,6 @@ def visualize_embeddings(
             # Log the error
             logger.error(f"Error during UMAP transformation: {e}")
             return []
+
+    def get_chunk_contents(self, chunk_id: str) -> BaseNode :
+        return self.llama_vector_store().get_nodes([chunk_id])[0]

llm-service/app/config.py

@@ -51,6 +51,7 @@
 
 SummaryStorageProviderType = Literal["Local", "S3"]
 ChatStoreProviderType = Literal["Local", "S3"]
+VectorDbProviderType = Literal["QDRANT", "OPENSEARCH"]
 
 
 class _Settings:
@@ -108,6 +109,18 @@ def vector_db_provider(self) -> Optional[str]:
     def opensearch_endpoint(self) -> str:
         return os.environ.get("OPENSEARCH_ENDPOINT", "http://localhost:9200")
 
+    @property
+    def opensearch_namespace(self) -> str:
+        return os.environ.get("OPENSEARCH_NAMESPACE", "rag_document_index")
+
+    @property
+    def opensearch_username(self) -> str:
+        return os.environ.get("OPENSEARCH_USERNAME", "")
+
+    @property
+    def opensearch_password(self) -> str:
+        return os.environ.get("OPENSEARCH_PASSWORD", "")
+
     @property
     def document_bucket_prefix(self) -> str:
         return os.environ.get("S3_RAG_BUCKET_PREFIX", "")

llm-service/app/routers/index/data_source/__init__.py

@@ -112,7 +112,7 @@ def delete(self, data_source_id: int) -> None:
     )
     @exceptions.propagates
     def chunk_contents(self, chunk_id: str) -> ChunkContentsResponse:
-        node = self.chunks_vector_store.llama_vector_store().get_nodes([chunk_id])[0]
+        node = self.chunks_vector_store.get_chunk_contents(chunk_id)
         return ChunkContentsResponse(
             text=node.get_content(),
             metadata=node.metadata,

llm-service/app/services/amp_metadata/__init__.py

@@ -40,9 +40,14 @@
 import socket
 from typing import Optional, cast, Protocol
 
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 
-from app.config import settings, SummaryStorageProviderType, ChatStoreProviderType
+from app.config import (
+    settings,
+    SummaryStorageProviderType,
+    ChatStoreProviderType,
+    VectorDbProviderType,
+)
 
 
 class AwsConfig(BaseModel):
@@ -83,6 +88,15 @@ class OpenAiConfig(BaseModel):
     openai_api_key: Optional[str] = None
     openai_api_base: Optional[str] = None
 
+
+class OpenSearchConfig(BaseModel):
+
+    opensearch_username: Optional[str] = None
+    opensearch_password: Optional[str] = None
+    opensearch_endpoint: Optional[str] = None
+    opensearch_namespace: Optional[str] = None
+
+
 class ProjectConfig(BaseModel):
     """
     Model to represent the project configuration.
@@ -91,10 +105,12 @@ class ProjectConfig(BaseModel):
     use_enhanced_pdf_processing: Optional[bool] = False
     summary_storage_provider: SummaryStorageProviderType
     chat_store_provider: ChatStoreProviderType
+    vector_db_provider: VectorDbProviderType
     aws_config: AwsConfig
     azure_config: AzureConfig
     caii_config: CaiiConfig
     openai_config: OpenAiConfig
+    opensearch_config: OpenSearchConfig
     cdp_token: Optional[str] = None
 
 
@@ -200,6 +216,7 @@ def config_to_env(config: ProjectConfig) -> dict[str, str]:
         "USE_ENHANCED_PDF_PROCESSING": str(config.use_enhanced_pdf_processing).lower(),
         "SUMMARY_STORAGE_PROVIDER": config.summary_storage_provider or "Local",
         "CHAT_STORE_PROVIDER": config.chat_store_provider or "Local",
+        "VECTOR_DB_PROVIDER": config.vector_db_provider or "QDRANT",
         "AWS_DEFAULT_REGION": config.aws_config.region or "",
         "S3_RAG_DOCUMENT_BUCKET": config.aws_config.document_bucket_name or "",
         "S3_RAG_BUCKET_PREFIX": config.aws_config.bucket_prefix or "",
@@ -209,6 +226,10 @@ def config_to_env(config: ProjectConfig) -> dict[str, str]:
         "AZURE_OPENAI_ENDPOINT": config.azure_config.openai_endpoint or "",
         "OPENAI_API_VERSION": config.azure_config.openai_api_version or "",
         "CAII_DOMAIN": config.caii_config.caii_domain or "",
+        "OPENSEARCH_USERNAME": config.opensearch_config.opensearch_username or "",
+        "OPENSEARCH_PASSWORD": config.opensearch_config.opensearch_password or "",
+        "OPENSEARCH_ENDPOINT": config.opensearch_config.opensearch_endpoint or "",
+        "OPENSEARCH_NAMESPACE": config.opensearch_config.opensearch_namespace or "",
         "OPENAI_API_KEY": config.openai_config.openai_api_key or "",
         "OPENAI_API_BASE": config.openai_config.openai_api_base or "",
     }
@@ -235,22 +256,33 @@ def build_configuration(
     caii_config = CaiiConfig(
         caii_domain=env.get("CAII_DOMAIN"),
     )
+    opensearch_config = OpenSearchConfig(
+        opensearch_username=env.get(
+            "OPENSEARCH_USERNAME",
+        ),
+        opensearch_password=env.get("OPENSEARCH_PASSWORD"),
+        opensearch_endpoint=env.get("OPENSEARCH_ENDPOINT"),
+        opensearch_namespace=env.get("OPENSEARCH_NAMESPACE"),
+    )
     return ProjectConfigPlus(
-        use_enhanced_pdf_processing=cast(
-            bool,
+        use_enhanced_pdf_processing=TypeAdapter(bool).validate_python(
             env.get("USE_ENHANCED_PDF_PROCESSING", False),
         ),
-        summary_storage_provider=cast(
-            SummaryStorageProviderType,
+        summary_storage_provider=TypeAdapter(
+            SummaryStorageProviderType
+        ).validate_python(
             env.get("SUMMARY_STORAGE_PROVIDER", "Local"),
         ),
-        chat_store_provider=cast(
-            ChatStoreProviderType,
+        chat_store_provider=TypeAdapter(ChatStoreProviderType).validate_python(
             env.get("CHAT_STORE_PROVIDER", "Local"),
         ),
+        vector_db_provider=TypeAdapter(VectorDbProviderType).validate_python(
+            env.get("VECTOR_DB_PROVIDER", "QDRANT")
+        ),
         aws_config=aws_config,
         azure_config=azure_config,
         caii_config=caii_config,
+        opensearch_config=opensearch_config,
         is_valid_config=validate(env),
         release_version=os.environ.get("RELEASE_TAG", "unknown"),
         application_config=application_config,