Enable kb-free chat sessions (#156)

* update python and java to allow for no kb

* "start enabling no-kb chats"

* wip on enabling ui to allow for chatting against no kb

* "wip ui"

* "wip on no-kb-chats"

* "add kb form item to chat settings"

* "formatting nodes wip"

* "add dsid to node"

* "invalidate suggested questions on session update"

* "new prompt for suggested q"

* "do suggested questions on a no-kb chat"

* fix azure env vars

* minor ui changes, adding claude 3.7

* remove claude 3.7

* "wip on empty chat cleanup"

* remove errant `bin` directory

* wip on fixing tests

* fix broken ui tests

* simplify mocking

* move vi mocks up to the top level so we don't fool ourselves into in thinking they are scoped

---------

Co-authored-by: Elijah Williams <[email protected]>

Author: jkwatson

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

@@ -194,7 +194,7 @@ public void delete(Long id) {
   public void update(Types.Session input) {
     var updatedInput = input.withTimeUpdated(Instant.now());
     String json = serializeQueryConfiguration(input);
-    jdbi.useHandle(
+    jdbi.useTransaction(
         handle -> {
           var sql =
               """
@@ -208,6 +208,11 @@ public void update(Types.Session input) {
               .bind("queryConfiguration", json)
               .bindMethods(updatedInput)
               .execute();
+          handle
+              .createUpdate("DELETE FROM CHAT_SESSION_DATA_SOURCE WHERE CHAT_SESSION_ID = :id")
+              .bind("id", input.id())
+              .execute();
+          insertSessionDataSources(handle, input.id(), input.dataSourceIds());
         });
   }
 

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

@@ -61,9 +61,14 @@ public static Types.Session createTestSessionInstance(String sessionName) {
   }
 
   public static Types.CreateSession createSessionInstance(String sessionName) {
+    return createSessionInstance(sessionName, List.of(1L, 2L, 3L));
+  }
+
+  public static Types.CreateSession createSessionInstance(
+      String sessionName, List<Long> dataSourceIds) {
     return new Types.CreateSession(
         sessionName,
-        List.of(1L, 2L, 3L),
+        dataSourceIds,
         "test-model",
         "test-rerank-model",
         3,

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

@@ -46,6 +46,7 @@
 import com.cloudera.cai.rag.util.UserTokenCookieDecoderTest;
 import com.cloudera.cai.util.exceptions.NotFound;
 import com.fasterxml.jackson.core.JsonProcessingException;
+import java.util.List;
 import org.junit.jupiter.api.Test;
 import org.springframework.mock.web.MockCookie;
 import org.springframework.mock.web.MockHttpServletRequest;
@@ -74,6 +75,29 @@ void create() throws JsonProcessingException {
     assertThat(result.queryConfiguration()).isNotNull();
   }
 
+  @Test
+  void create_noDataSource() throws JsonProcessingException {
+    SessionController sessionController = new SessionController(SessionService.createNull());
+    var request = new MockHttpServletRequest();
+    request.setCookies(
+        new MockCookie("_basusertoken", UserTokenCookieDecoderTest.encodeCookie("test-user")));
+    var sessionName = "test";
+    Types.CreateSession input = TestData.createSessionInstance(sessionName, List.of());
+    Types.Session result = sessionController.create(input, request);
+    assertThat(result.id()).isNotNull();
+    assertThat(result.name()).isEqualTo(sessionName);
+    assertThat(result.inferenceModel()).isEqualTo(input.inferenceModel());
+    assertThat(result.rerankModel()).isEqualTo(input.rerankModel());
+    assertThat(result.responseChunks()).isEqualTo(input.responseChunks());
+    assertThat(result.dataSourceIds()).isEmpty();
+    assertThat(result.timeCreated()).isNotNull();
+    assertThat(result.timeUpdated()).isNotNull();
+    assertThat(result.createdById()).isEqualTo("test-user");
+    assertThat(result.updatedById()).isEqualTo("test-user");
+    assertThat(result.lastInteractionTime()).isNull();
+    assertThat(result.queryConfiguration()).isNotNull();
+  }
+
   @Test
   void get() {
     SessionController sessionController = new SessionController(SessionService.createNull());

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

@@ -42,6 +42,7 @@
 
 import com.cloudera.cai.rag.TestData;
 import com.cloudera.cai.rag.Types;
+import java.util.List;
 import org.junit.jupiter.api.Test;
 
 class SessionServiceTest {
@@ -77,9 +78,10 @@ void update() {
             TestData.createTestSessionInstance("test")
                 .withCreatedById("abc")
                 .withUpdatedById("abc"));
-    var updated = result.withRerankModel("");
+    var updated = result.withRerankModel("").withDataSourceIds(List.of(4L));
     var updatedResult = sessionService.update(updated);
     assertThat(updatedResult.rerankModel()).isNull();
+    assertThat(updatedResult.dataSourceIds()).containsExactly(4L);
   }
 
   @Test

cacerts.jks

@@ -47,6 +47,7 @@
 from ....rag_types import RagPredictConfiguration
 from ....services.chat import generate_suggested_questions, v2_chat, direct_llm_chat
 from ....services.chat_store import ChatHistoryManager, RagStudioChatMessage
+from ....services.metadata_apis import session_metadata_api
 from ....services.mlflow import rating_mlflow_log_metric, feedback_mlflow_log_table
 
 logger = logging.getLogger(__name__)
@@ -141,11 +142,12 @@ def chat(
     _basusertoken: Annotated[str | None, Cookie()] = None,
 ) -> RagStudioChatMessage:
     user_name = parse_jwt_cookie(_basusertoken)
+    session = session_metadata_api.get_session(session_id)
 
     configuration = request.configuration or RagPredictConfiguration()
-    if configuration.exclude_knowledge_base:
-        return direct_llm_chat(session_id, request.query, user_name)
-    return v2_chat(session_id, request.query, configuration, user_name)
+    if configuration.exclude_knowledge_base or len(session.data_source_ids) == 0:
+        return direct_llm_chat(session, request.query, user_name)
+    return v2_chat(session, request.query, configuration, user_name)
 
 
 class RagSuggestedQuestionsResponse(BaseModel):

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

@@ -62,9 +62,8 @@
 
 
 def v2_chat(
-    session_id: int, query: str, configuration: RagPredictConfiguration, user_name: str
+    session: Session, query: str, configuration: RagPredictConfiguration, user_name: str
 ) -> RagStudioChatMessage:
-    session = session_metadata_api.get_session(session_id)
     query_configuration = QueryConfiguration(
         top_k=session.response_chunks,
         model_name=session.inference_model,
@@ -80,7 +79,7 @@ def v2_chat(
         session, response_id, query, query_configuration, user_name
     )
 
-    ChatHistoryManager().append_to_history(session_id, [new_chat_message])
+    ChatHistoryManager().append_to_history(session.id, [new_chat_message])
     return new_chat_message
 
 
@@ -121,7 +120,7 @@ def _run_chat(
     relevance, faithfulness = evaluators.evaluate_response(
         query, response, session.inference_model
     )
-    response_source_nodes = format_source_nodes(response)
+    response_source_nodes = format_source_nodes(response, data_source_id)
     new_chat_message = RagStudioChatMessage(
         id=response_id,
         source_nodes=response_source_nodes,
@@ -159,7 +158,9 @@ def retrieve_chat_history(session_id: int) -> List[RagContext]:
     return history
 
 
-def format_source_nodes(response: AgentChatResponse) -> List[RagPredictSourceNode]:
+def format_source_nodes(
+    response: AgentChatResponse, data_source_id: int
+) -> List[RagPredictSourceNode]:
     response_source_nodes = []
     for source_node in response.source_nodes:
         doc_id = source_node.node.metadata.get("document_id", source_node.node.node_id)
@@ -169,6 +170,7 @@ def format_source_nodes(response: AgentChatResponse) -> List[RagPredictSourceNod
                 doc_id=doc_id,
                 source_file_name=source_node.node.metadata["file_name"],
                 score=source_node.score or 0.0,
+                dataSourceId=data_source_id,
             )
         )
     response_source_nodes = sorted(
@@ -177,10 +179,32 @@ def format_source_nodes(response: AgentChatResponse) -> List[RagPredictSourceNod
     return response_source_nodes
 
 
+def generate_suggested_questions_direct_llm(session: Session) -> List[str]:
+    chat_history = retrieve_chat_history(session.id)
+    if not chat_history:
+        return []
+    query_str = (
+        " Give me a list of possible follow-up questions."
+        " Each question should be on a new line."
+        " There should be no more than four (4) questions."
+        " Each question should be no longer than fifteen (15) words."
+        " The response should be a bulleted list, using an asterisk (*) to denote the bullet item."
+        " Do not start like this - `Here are four questions that I can answer based on the context information`"
+        " Only return the list."
+    )
+    chat_response = llm_completion.completion(
+        session.id, query_str, session.inference_model
+    )
+    suggested_questions = process_response(chat_response.message.content)
+    return suggested_questions
+
+
 def generate_suggested_questions(
     session_id: int,
 ) -> List[str]:
     session = session_metadata_api.get_session(session_id)
+    if len(session.data_source_ids) == 0:
+        return generate_suggested_questions_direct_llm(session)
     if len(session.data_source_ids) != 1:
         raise HTTPException(
             status_code=400,
@@ -256,14 +280,13 @@ def process_response(response: str | None) -> list[str]:
 
 
 def direct_llm_chat(
-    session_id: int, query: str, user_name: str
+    session: Session, query: str, user_name: str
 ) -> RagStudioChatMessage:
-    session = session_metadata_api.get_session(session_id)
     response_id = str(uuid.uuid4())
     record_direct_llm_mlflow_run(response_id, session, user_name)
 
     chat_response = llm_completion.completion(
-        session_id, query, session.inference_model
+        session.id, query, session.inference_model
     )
     new_chat_message = RagStudioChatMessage(
         id=response_id,
@@ -277,5 +300,5 @@ def direct_llm_chat(
         timestamp=time.time(),
         condensed_question=None,
     )
-    ChatHistoryManager().append_to_history(session_id, [new_chat_message])
+    ChatHistoryManager().append_to_history(session.id, [new_chat_message])
     return new_chat_message

llm-service/app/services/chat.py

@@ -51,6 +51,7 @@ class RagPredictSourceNode(BaseModel):
     doc_id: str
     source_file_name: str
     score: float
+    dataSourceId: Optional[int] = None
 
 
 class Evaluation(BaseModel):
@@ -119,7 +120,7 @@ def retrieve_chat_history(self, session_id: int) -> List[RagStudioChatMessage]:
                         "evaluations", []
                     ),
                     timestamp=assistant_message.additional_kwargs.get("timestamp", 0.0),
-                    condensed_question=None
+                    condensed_question=None,
                 )
             )
             i += 2

llm-service/app/services/chat_store.py

@@ -66,13 +66,13 @@
 from ..query.simple_reranker import SimpleReranker
 
 __all__ = [
-    'CAIIModelProvider',
-    'ModelType',
-    'Embedding',
-    'LLM',
-    'Reranking',
-    'ModelSource',
-    'BedrockModelProvider'
+    "CAIIModelProvider",
+    "ModelType",
+    "Embedding",
+    "LLM",
+    "Reranking",
+    "ModelSource",
+    "BedrockModelProvider",
 ]
 
 T = TypeVar("T", bound=BaseComponent)

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

@@ -45,8 +45,7 @@
 class AzureModelProvider(ModelProvider):
     @staticmethod
     def get_env_var_names() -> set[str]:
-        return {"AZURE_OPENAI_API_KEY" "AZURE_OPENAI_ENDPOINT" "OPENAI_API_VERSION"}
-
+        return {"AZURE_OPENAI_API_KEY", "AZURE_OPENAI_ENDPOINT", "OPENAI_API_VERSION"}
     @staticmethod
     def get_llm_models() -> List[ModelResponse]:
         return [