Merge pull request #311 from cloudera/main

Release 1.29.0

Author: ewilliams-cloudera

backend/src/main/java/com/cloudera/cai/rag/Types.java

@@ -103,10 +103,12 @@ public record RagDataSource(
       Long associatedSessionId) {}
 
   @With
+  @Builder
   public record QueryConfiguration(
       boolean enableHyde,
       boolean enableSummaryFilter,
       boolean enableToolCalling,
+      Boolean disableStreaming,
       List<String> selectedTools) {}
 
   @With

backend/src/main/java/com/cloudera/cai/rag/sessions/SessionRepository.java

@@ -43,6 +43,7 @@
 import com.cloudera.cai.rag.configuration.JdbiConfiguration;
 import com.cloudera.cai.util.exceptions.NotFound;
 import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import java.time.Instant;
 import java.util.List;
@@ -57,9 +58,16 @@
 @Component
 public class SessionRepository {
   public static final Types.QueryConfiguration DEFAULT_QUERY_CONFIGURATION =
-      new Types.QueryConfiguration(false, true, false, List.of());
+      Types.QueryConfiguration.builder()
+          .enableHyde(false)
+          .enableSummaryFilter(true)
+          .enableToolCalling(false)
+          .disableStreaming(true)
+          .selectedTools(List.of())
+          .build();
   private final DatabaseOperations databaseOperations;
-  private final ObjectMapper objectMapper = new ObjectMapper();
+  private final ObjectMapper objectMapper =
+      new ObjectMapper().disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES);
 
   public SessionRepository(DatabaseOperations databaseOperations) {
     this.databaseOperations = databaseOperations;
@@ -177,6 +185,9 @@ private Types.QueryConfiguration extractQueryConfiguration(RowView rowView)
     if (queryConfiguration.selectedTools() == null) {
       queryConfiguration = queryConfiguration.withSelectedTools(List.of());
     }
+    if (queryConfiguration.disableStreaming() == null) {
+      queryConfiguration = queryConfiguration.withDisableStreaming(false);
+    }
     return queryConfiguration;
   }
 

backend/src/test/java/com/cloudera/cai/rag/TestData.java

@@ -83,7 +83,12 @@ public static Types.Session createTestSessionInstance(
         null,
         "test-rerank-model",
         3,
-        new Types.QueryConfiguration(false, true, true, List.of()));
+        Types.QueryConfiguration.builder()
+            .enableSummaryFilter(true)
+            .enableToolCalling(true)
+            .disableStreaming(true)
+            .selectedTools(List.of())
+            .build());
   }
 
   public static Types.CreateSession createSessionInstance(String sessionName) {
@@ -99,7 +104,13 @@ public static Types.CreateSession createSessionInstance(
         null,
         "test-rerank-model",
         3,
-        new Types.QueryConfiguration(false, true, true, List.of()),
+        Types.QueryConfiguration.builder()
+            .enableHyde(false)
+            .enableSummaryFilter(true)
+            .enableToolCalling(true)
+            .disableStreaming(true)
+            .selectedTools(List.of())
+            .build(),
         projectId);
   }
 

backend/src/test/java/com/cloudera/cai/rag/sessions/SessionControllerTest.java

@@ -137,6 +137,14 @@ void update() {
     var updatedRerankModel = "new-rerank-model";
     var updatedProjectId = project.id();
 
+    var queryConfiguration =
+        Types.QueryConfiguration.builder()
+            .enableHyde(true)
+            .enableSummaryFilter(false)
+            .enableToolCalling(true)
+            .disableStreaming(true)
+            .selectedTools(List.of("foo"))
+            .build();
     var updatedSession =
         sessionController.update(
             insertedSession
@@ -145,8 +153,7 @@ void update() {
                 .withRerankModel(updatedRerankModel)
                 .withName(updatedName)
                 .withProjectId(updatedProjectId)
-                .withQueryConfiguration(
-                    new Types.QueryConfiguration(true, false, true, List.of("foo"))),
+                .withQueryConfiguration(queryConfiguration),
             request);
 
     assertThat(updatedSession.id()).isNotNull();
@@ -160,8 +167,7 @@ void update() {
     assertThat(updatedSession.timeUpdated()).isAfter(insertedSession.timeUpdated());
     assertThat(updatedSession.createdById()).isEqualTo("test-user");
     assertThat(updatedSession.lastInteractionTime()).isNull();
-    assertThat(updatedSession.queryConfiguration())
-        .isEqualTo(new Types.QueryConfiguration(true, false, true, List.of("foo")));
+    assertThat(updatedSession.queryConfiguration()).isEqualTo(queryConfiguration);
   }
 
   @Test

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

@@ -112,7 +112,11 @@ def _find_dim(data_source_id: int) -> int:
 
     def size(self) -> Optional[int]:
         os_client = self._low_level_client
-        return int(os_client.count(index=self.table_name)["count"])
+        try:
+            return int(os_client.count(index=self.table_name)["count"])
+        except opensearchpy.exceptions.NotFoundError:
+            # Return 0 if index doesn't exist yet
+            return 0
 
     def delete(self) -> None:
         os_client = self._low_level_client

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

@@ -301,6 +301,18 @@ def generate_stream() -> Generator[str, None, None]:
 
             first_message = True
             stream = future.result()
+
+            # If streaming is disabled, immediately send a loading event to show StreamedEvents
+            if session.query_configuration.disable_streaming:
+                loading = ChatEvent(
+                    type="thinking",
+                    name="thinking",
+                    timestamp=time.time(),
+                    data="Preparing full response...",
+                )
+                event_json = json.dumps({"event": loading.model_dump()})
+                yield f"data: {event_json}\n\n"
+                first_message = False
             for item in stream:
                 response: ChatResponse = item
                 # Check for cancellation between each response

llm-service/app/services/caii/caii.py

@@ -302,9 +302,11 @@ def get_embedding_model(model_name: str) -> BaseEmbedding:
 
 
 def get_caii_llm_models() -> List[ModelResponse]:
-    potential_models = get_models_with_task("TEXT_GENERATION")
+    potential_text_models = get_models_with_task("TEXT_GENERATION")
+    potential_text_to_text_models = get_models_with_task("TEXT_TO_TEXT_GENERATION")
+
     results: list[Endpoint] = []
-    for potential in potential_models:
+    for potential in [*potential_text_models, *potential_text_to_text_models]:
         try:
             model = get_llm(
                 endpoint=potential,

llm-service/app/services/chat/streaming_chat.py

@@ -83,6 +83,7 @@ def stream_chat(
         use_hyde=session.query_configuration.enable_hyde,
         use_summary_filter=session.query_configuration.enable_summary_filter,
         use_tool_calling=session.query_configuration.enable_tool_calling,
+        use_streaming=not session.query_configuration.disable_streaming,
     )
 
     response_id = str(uuid.uuid4())
@@ -188,14 +189,20 @@ def _stream_direct_llm_chat(
     user_name: Optional[str],
 ) -> Generator[ChatResponse, None, None]:
     record_direct_llm_mlflow_run(response_id, session, user_name)
-
-    chat_response = llm_completion.stream_completion(
-        session.id, query, session.inference_model
-    )
-    response: ChatResponse = ChatResponse(message=ChatMessage(content=query))
-    for response in chat_response:
+    response: ChatResponse
+    if session.query_configuration.disable_streaming:
+        # Use non-streaming completion when streaming is disabled
+        response = llm_completion.completion(session.id, query, session.inference_model)
         response.additional_kwargs["response_id"] = response_id
         yield response
+    else:
+        chat_response = llm_completion.stream_completion(
+            session.id, query, session.inference_model
+        )
+        response = ChatResponse(message=ChatMessage(content=query))
+        for response in chat_response:
+            response.additional_kwargs["response_id"] = response_id
+            yield response
 
     new_chat_message = RagStudioChatMessage(
         id=response_id,

llm-service/app/services/metadata_apis/session_metadata_api.py

@@ -52,6 +52,7 @@ class SessionQueryConfiguration:
     enable_summary_filter: bool
     enable_tool_calling: bool = False
     selected_tools: list[str] = field(default_factory=list)
+    disable_streaming: bool = False
 
 
 @dataclass
@@ -126,6 +127,9 @@ def session_from_java_response(data: dict[str, Any]) -> Session:
             enable_tool_calling=data["queryConfiguration"].get(
                 "enableToolCalling", False
             ),
+            disable_streaming=data["queryConfiguration"].get(
+                "disableStreaming", False
+            ),
             selected_tools=data["queryConfiguration"]["selectedTools"] or [],
         ),
         associated_data_source_id=data.get("associatedDataSourceId", None),
@@ -146,6 +150,7 @@ def update_session(session: Session, user_name: Optional[str]) -> Session:
             "enableSummaryFilter": session.query_configuration.enable_summary_filter,
             "enableToolCalling": session.query_configuration.enable_tool_calling,
             "selectedTools": session.query_configuration.selected_tools,
+            "disableStreaming": session.query_configuration.disable_streaming,
         },
         associatedDataSourceId=session.associated_data_source_id,
     )