add selectedTools as an attribute of the session queryConfiguration

Author: jkwatson

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

@@ -101,8 +101,12 @@ public record RagDataSource(
       @Nullable Long totalDocSize,
       boolean availableForDefaultProject) {}
 
+  @With
   public record QueryConfiguration(
-      boolean enableHyde, boolean enableSummaryFilter, boolean enableToolCalling) {}
+      boolean enableHyde,
+      boolean enableSummaryFilter,
+      boolean enableToolCalling,
+      List<String> selectedTools) {}
 
   @With
   @Builder

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

@@ -57,7 +57,7 @@
 @Component
 public class SessionRepository {
   public static final Types.QueryConfiguration DEFAULT_QUERY_CONFIGURATION =
-      new Types.QueryConfiguration(false, true, false);
+      new Types.QueryConfiguration(false, true, false, List.of());
   private final Jdbi jdbi;
   private final ObjectMapper objectMapper = new ObjectMapper();
 
@@ -169,6 +169,9 @@ private Types.QueryConfiguration extractQueryConfiguration(RowView rowView)
     if (queryConfiguration == null) {
       return DEFAULT_QUERY_CONFIGURATION;
     }
+    if (queryConfiguration.selectedTools() == null) {
+      queryConfiguration = queryConfiguration.withSelectedTools(List.of());
+    }
     return queryConfiguration;
   }
 

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

@@ -82,7 +82,7 @@ public static Types.Session createTestSessionInstance(
         "test-model",
         "test-rerank-model",
         3,
-        new Types.QueryConfiguration(false, true, true));
+        new Types.QueryConfiguration(false, true, true, List.of()));
   }
 
   public static Types.CreateSession createSessionInstance(String sessionName) {
@@ -97,7 +97,7 @@ public static Types.CreateSession createSessionInstance(
         "test-model",
         "test-rerank-model",
         3,
-        new Types.QueryConfiguration(false, true, true),
+        new Types.QueryConfiguration(false, true, true, List.of()),
         projectId);
   }
 

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

@@ -145,7 +145,8 @@ void update() {
                 .withRerankModel(updatedRerankModel)
                 .withName(updatedName)
                 .withProjectId(updatedProjectId)
-                .withQueryConfiguration(new Types.QueryConfiguration(true, false, true)),
+                .withQueryConfiguration(
+                    new Types.QueryConfiguration(true, false, true, List.of("foo"))),
             request);
 
     assertThat(updatedSession.id()).isNotNull();
@@ -160,7 +161,7 @@ void update() {
     assertThat(updatedSession.createdById()).isEqualTo("test-user");
     assertThat(updatedSession.lastInteractionTime()).isNull();
     assertThat(updatedSession.queryConfiguration())
-        .isEqualTo(new Types.QueryConfiguration(true, false, true));
+        .isEqualTo(new Types.QueryConfiguration(true, false, true, List.of("foo")));
   }
 
   @Test

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

@@ -36,7 +36,7 @@
 #  DATA.
 #
 import json
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import datetime
 from typing import List, Any, Optional
 
@@ -51,6 +51,7 @@ class SessionQueryConfiguration:
     enable_hyde: bool
     enable_summary_filter: bool
     enable_tool_calling: bool = False
+    selected_tools: list[str] = field(default_factory=list)
 
 
 @dataclass
@@ -116,6 +117,7 @@ def session_from_java_response(data: dict[str, Any]) -> Session:
             enable_tool_calling=data["queryConfiguration"].get(
                 "enableToolCalling", False
             ),
+            selected_tools=data["queryConfiguration"]["selectedTools"] or []
         ),
     )
 
@@ -133,6 +135,7 @@ def update_session(session: Session, user_name: Optional[str]) -> Session:
             "enableHyde": session.query_configuration.enable_hyde,
             "enableSummaryFilter": session.query_configuration.enable_summary_filter,
             "enableToolCalling": session.query_configuration.enable_tool_calling,
+            "selectedTools": session.query_configuration.selected_tools
         },
     )
     headers = {