Mob/main (#287)

* allow an optional env var for the postgres admin role

* don't always check for CAII enablement

* refactor: optimize streaming chat mutation handlers with useCallback and useMemo

* Revert "refactor: optimize streaming chat mutation handlers with useCallback and useMemo"

This reverts commit 691b10f.

* potential solution for queueing chunks

* custom hook for buffering chunks

* fix summary reconciler

* bump mui packages

* remove comment in test

* rename file as no longer private

* fix imports

* remove import

---------

Co-authored-by: jwatson <[email protected]>
Co-authored-by: Baasit Sharief <[email protected]>

Author: 3 people

backend/src/main/java/com/cloudera/cai/rag/configuration/JdbiConfiguration.java

@@ -43,6 +43,7 @@
 import com.cloudera.cai.util.db.RdbConfig;
 import com.cloudera.cai.util.db.migration.Migrator;
 import java.sql.SQLException;
+import java.util.Optional;
 import javax.sql.DataSource;
 import lombok.extern.slf4j.Slf4j;
 import org.jdbi.v3.core.HandleCallback;
@@ -105,7 +106,11 @@ private static DatabaseConfig createDatabaseConfig() {
             .build();
     if (rdbConfiguration.isPostgres()) {
       rdbConfiguration =
-          rdbConfiguration.toBuilder().rdbUsername("postgres").rdbDatabaseName(null).build();
+          rdbConfiguration.toBuilder()
+              .rdbUsername(
+                  Optional.ofNullable(System.getenv("POSTGRES_ADMIN_ROLE")).orElse("postgres"))
+              .rdbDatabaseName(null)
+              .build();
     }
     return DatabaseConfig.builder().RdbConfiguration(rdbConfiguration).build();
   }

backend/src/main/java/com/cloudera/cai/rag/files/RagFileService.java

@@ -151,7 +151,7 @@ private RagDocumentMetadata processFile(
     log.info("Saved document with id: {}", id);
 
     ragFileIndexReconciler.submit(ragDocument.withId(id));
-    ragFileSummaryReconciler.submit(ragDocument.withId(id));
+    ragFileSummaryReconciler.resync();
 
     return new RagDocumentMetadata(
         ragDocument.filename(), documentId, ragDocument.extension(), ragDocument.sizeInBytes());

backend/src/main/java/com/cloudera/cai/rag/files/RagFileSummaryReconciler.java

@@ -91,7 +91,6 @@ public void resync() {
         SELECT rdsd.* from rag_data_source_document rdsd
          JOIN rag_data_source rds ON rdsd.data_source_id = rds.id
          WHERE rdsd.summary_creation_timestamp IS NULL
-           AND rdsd.VECTOR_UPLOAD_TIMESTAMP IS NOT NULL
            AND (rdsd.deleted IS NULL OR rdsd.deleted = :deleted)
            AND (rdsd.time_created > :yesterday OR rds.time_updated > :yesterday)
            AND rds.summarization_model IS NOT NULL AND rds.summarization_model != ''

llm-service/app/ai/indexing/summary_indexer.py

@@ -75,7 +75,8 @@
 from ..vector_stores.vector_store_factory import VectorStoreFactory
 from ...config import settings
 from ...services.metadata_apis import data_sources_metadata_api
-from ...services.models.providers import CAIIModelProvider, AzureModelProvider, OpenAiModelProvider
+from ...services.models.providers import ModelProvider
+from ...services.models import ModelSource
 
 logger = logging.getLogger(__name__)
 
@@ -132,11 +133,16 @@ def __index_configuration(
         embed_summaries: bool = True,
     ) -> Dict[str, Any]:
         prompt_helper: Optional[PromptHelper] = None
-        # if we're using CAII, let's be conservative and use a small context window to account for mistral's small context
-        if CAIIModelProvider.is_enabled():
+        model_source: ModelSource = (
+            ModelProvider.get_provider_class().get_model_source()
+        )
+        if model_source == "CAII":
+            # if we're using CAII, let's be conservative and use a small context window to account for mistral's small context
             prompt_helper = PromptHelper(context_window=3000)
-        if AzureModelProvider.is_enabled() or OpenAiModelProvider.is_enabled():
-            prompt_helper = PromptHelper(context_window=min(llm.metadata.context_window, 10000))
+        if model_source == "Azure" or model_source == "OpenAI":
+            prompt_helper = PromptHelper(
+                context_window=min(llm.metadata.context_window, 10000)
+            )
         return {
             "llm": llm,
             "response_synthesizer": get_response_synthesizer(

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

@@ -58,6 +58,7 @@
     AzureModelProvider,
     BedrockModelProvider,
 )
+from app.services.utils import timed_lru_cache
 
 
 class AwsConfig(BaseModel):
@@ -286,7 +287,9 @@ def validate_model_config(environ: dict[str, str]) -> ValidationResult:
     return ValidationResult(valid=valid_model_config_exists, message=message)
 
 
-def validate(environ: dict[str, str]) -> ConfigValidationResults:
+@timed_lru_cache(maxsize=1, seconds=6000)
+def validate(frozen_env: frozenset[tuple[str, str]]) -> ConfigValidationResults:
+    environ = {k: v for k, v in frozen_env}
     print("Validating environment variables...")
     storage_config = validate_storage_config(environ)
     model_config = validate_model_config(environ)
@@ -382,7 +385,7 @@ def build_configuration(
         opensearch_endpoint=env.get("OPENSEARCH_ENDPOINT"),
         opensearch_namespace=env.get("OPENSEARCH_NAMESPACE"),
     )
-    validate_config = validate(env)
+    validate_config = validate(frozenset(env.items()))
 
     model_provider = (
         TypeAdapter(ModelProviderType).validate_python(env.get("MODEL_PROVIDER"))

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

@@ -35,13 +35,13 @@
 #  BUSINESS ADVANTAGE OR UNAVAILABILITY, OR LOSS OR CORRUPTION OF
 #  DATA.
 #
-from ._model_source import ModelSource
 from .embedding import Embedding
 from .llm import LLM
-from .providers._model_provider import ModelProvider
+from .providers import ModelProvider
 from .reranking import Reranking
+from ._model_source import ModelSource
 
-__all__ = ["Embedding", "LLM", "Reranking"]
+__all__ = ["Embedding", "LLM", "Reranking", "ModelSource"]
 
 
 def get_model_source() -> ModelSource:

llm-service/app/services/models/providers/__init__.py

@@ -39,10 +39,12 @@
 from .bedrock import BedrockModelProvider
 from .caii import CAIIModelProvider
 from .openai import OpenAiModelProvider
+from ._model_provider import ModelProvider
 
 __all__ = [
     "AzureModelProvider",
     "BedrockModelProvider",
     "CAIIModelProvider",
     "OpenAiModelProvider",
-]
+    "ModelProvider",
+]

ui/package.json

@@ -23,11 +23,11 @@
     "@emotion/react": "^11.14.0",
     "@emotion/styled": "^11.14.0",
     "@microsoft/fetch-event-source": "^2.0.1",
-    "@mui/material": "^6.4.3",
-    "@mui/x-charts": "^7.26.0",
+    "@mui/material": "^7.3.1",
+    "@mui/x-charts": "^8.9.2",
     "@tanstack/react-query": "^5.79.0",
     "@tanstack/react-query-devtools": "^5.79.0",
-    "@tanstack/react-router": "^1.120.13",
+    "@tanstack/react-router": "^1.130.12",
     "antd": "^5.25.3",
     "date-fns": "^4.1.0",
     "lodash": "^4.17.21",