Chat UI improvements (#270)

* move select a datasource

* fix search

* fix spacing in input

* width

* add inference model selection to the chat input

* fix broken test

* handle showing correct model before activeSession

* handle tool calling enablement based on inference model

* adjust width of options

* fix up broken things, and provide models for projects to the chat initiator

* adjust colors

* capitalize

* bug fix for retriever tool with no summaries

* fix broken tests

* split bedrock by model family

---------

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

Author: 3 people

llm-service/app/services/query/agents/tool_calling_querier.py

@@ -111,7 +111,8 @@ def should_use_retrieval(
         data_source_summary_indexer = SummaryIndexer.get_summary_indexer(data_source_id)
         if data_source_summary_indexer:
             data_source_summary = data_source_summary_indexer.get_full_summary()
-            data_source_summaries[data_source_id] = data_source_summary
+            if data_source_summary:
+                data_source_summaries[data_source_id] = data_source_summary
     return len(data_source_ids) > 0, data_source_summaries
 
 

ui/src/api/modelsApi.ts

@@ -47,7 +47,7 @@ import {
 } from "src/api/utils.ts";
 
 export interface Model {
-  name?: string;
+  name: string;
   model_id: string;
   available: boolean | null;
   replica_count?: number;

ui/src/pages/Analytics/AnalyticsPage.tsx

@@ -45,7 +45,10 @@ import {
   Select,
   Typography,
 } from "antd";
-import { transformModelOptions } from "src/utils/modelUtils.ts";
+import {
+  ModelSelectOptions,
+  useTransformModelOptions,
+} from "src/utils/modelUtils.ts";
 import { useGetLlmModels, useGetRerankingModels } from "src/api/modelsApi.ts";
 import { MetricFilter } from "src/api/metricsApi.ts";
 import MetadataMetrics from "pages/Analytics/AppMetrics.tsx";
@@ -79,7 +82,7 @@ const SelectFilterOption = ({
 }: {
   name: string;
   label: string;
-  options: { value: string; label: string | undefined }[];
+  options: ModelSelectOptions;
 }) => {
   return (
     <Form.Item name={name} label={label} style={{ marginTop: 8 }}>
@@ -106,6 +109,8 @@ const MetricFilterOptions = ({
   const { data: llmModels } = useGetLlmModels();
   const { data: rerankingModels } = useGetRerankingModels();
   const { data: projects } = useGetProjects();
+  const llmModelOptions = useTransformModelOptions(llmModels);
+  const rerankModelOptions = useTransformModelOptions(rerankingModels);
 
   return (
     <Form
@@ -117,15 +122,12 @@ const MetricFilterOptions = ({
       <SelectFilterOption
         name="inference_model"
         label="Response synthesizer model"
-        options={transformModelOptions(llmModels)}
+        options={llmModelOptions}
       />
       <SelectFilterOption
         name="rerank_model"
         label="Reranking model"
-        options={[
-          ...transformModelOptions(rerankingModels),
-          { value: "none", label: "None" },
-        ]}
+        options={[...rerankModelOptions, { value: "none", label: "None" }]}
       />
       <BooleanFilterOption
         name="use_summary_filter"

ui/src/pages/DataSources/DataSourcesManagement/DataSourcesForm.tsx

@@ -52,7 +52,7 @@ import {
 import { ConnectionType, DataSourceBaseType } from "src/api/dataSourceApi";
 import { useGetEmbeddingModels, useGetLlmModels } from "src/api/modelsApi.ts";
 import { useEffect } from "react";
-import { transformModelOptions } from "src/utils/modelUtils.ts";
+import { useTransformModelOptions } from "src/utils/modelUtils.ts";
 import { useNavigate } from "@tanstack/react-router";
 import messageQueue from "src/utils/messageQueue.ts";
 
@@ -148,6 +148,8 @@ const DataSourcesForm = ({
   const embeddingsModels = useGetEmbeddingModels();
   const llmModels = useGetLlmModels();
   const navigate = useNavigate();
+  const embeddingModelOptions = useTransformModelOptions(embeddingsModels.data);
+  const llmModelOptions = useTransformModelOptions(llmModels.data);
 
   useEffect(() => {
     if (initialValues.embeddingModel) {
@@ -226,7 +228,7 @@ const DataSourcesForm = ({
         initialValue={initialValues.embeddingModel}
       >
         <Select
-          options={transformModelOptions(embeddingsModels.data)}
+          options={embeddingModelOptions}
           disabled={updateMode}
           loading={embeddingsModels.isLoading}
         />
@@ -246,7 +248,7 @@ const DataSourcesForm = ({
         initialValue={initialValues.summarizationModel}
       >
         <Select
-          options={transformModelOptions(llmModels.data)}
+          options={llmModelOptions}
           allowClear
           loading={llmModels.isLoading}
         />

ui/src/pages/Projects/ProjectPage/NewChatSession/NewChatSession.tsx

@@ -36,19 +36,18 @@
  * DATA.
  ******************************************************************************/
 
-import { Card, Select, Typography } from "antd";
-import RagChatQueryInput from "pages/RagChatTab/FooterComponents/RagChatQueryInput.tsx";
+import { Card, Typography } from "antd";
+import RagChatQueryInput, {
+  NewSessionCallbackProps,
+} from "pages/RagChatTab/FooterComponents/RagChatQueryInput.tsx";
 import useCreateSessionAndRedirect from "pages/RagChatTab/ChatOutput/hooks/useCreateSessionAndRedirect.tsx";
 import { useGetDataSourcesForProject } from "src/api/projectsApi.ts";
 import { useProjectContext } from "pages/Projects/ProjectContext.tsx";
-import { DataSourceInputType, formatDataSource } from "src/utils/formatters.ts";
-import { useState } from "react";
 
 export const NewChatSession = () => {
   const { project } = useProjectContext();
   const createSessionAndRedirect = useCreateSessionAndRedirect();
   const { data: dataSources } = useGetDataSourcesForProject(project.id);
-  const [selectedDataSources, setSelectedDataSources] = useState<number[]>([]);
 
   return (
     <Card
@@ -57,25 +56,19 @@ export const NewChatSession = () => {
           Start a new chat session
         </Typography.Title>
       }
-      extra={
-        <Select
-          mode="multiple"
-          placeholder="Select a data source (optional)"
-          disabled={!dataSources || dataSources.length === 0}
-          style={{ width: 300 }}
-          allowClear={true}
-          onChange={(ids: DataSourceInputType["value"][]) => {
-            setSelectedDataSources(ids);
-          }}
-          options={dataSources?.map((value) => {
-            return formatDataSource(value);
-          })}
-        />
-      }
     >
       <RagChatQueryInput
-        newSessionCallback={(userInput: string) => {
-          createSessionAndRedirect(selectedDataSources, userInput);
+        validDataSources={dataSources}
+        newSessionCallback={({
+          userInput,
+          selectedDataSourceIds,
+          inferenceModel,
+        }: NewSessionCallbackProps) => {
+          createSessionAndRedirect(
+            selectedDataSourceIds,
+            userInput,
+            inferenceModel,
+          );
         }}
       />
     </Card>

ui/src/pages/RagChatTab/ChatOutput/Placeholders/NoDataSourcesState.tsx

@@ -37,12 +37,9 @@
  ******************************************************************************/
 
 import { useNavigate } from "@tanstack/react-router";
-import { Button, Flex, Form, Select, Typography } from "antd";
+import { Button, Flex, Typography } from "antd";
 import { useContext, ReactNode } from "react";
 import { RagChatContext } from "pages/RagChatTab/State/RagChatContext.tsx";
-import useCreateSessionAndRedirect from "pages/RagChatTab/ChatOutput/hooks/useCreateSessionAndRedirect.tsx";
-import { ArrowRightOutlined } from "@ant-design/icons";
-import { formatDataSource } from "src/utils/formatters.ts";
 
 const PlaceholderContainer = ({
   message,
@@ -70,63 +67,21 @@ const PlaceholderContainer = ({
 
 const NoDataSourcesState = () => {
   const navigate = useNavigate();
-  const [form] = Form.useForm<{ dataSourceIds: number[] }>();
   const {
     activeSession,
     dataSourcesQuery: { dataSources, dataSourcesStatus },
   } = useContext(RagChatContext);
 
-  const createSessionAndRedirect = useCreateSessionAndRedirect();
-
-  const handleCreateSession = () => {
-    form
-      .validateFields()
-      .catch(() => null)
-      .then((values) => {
-        if (values?.dataSourceIds.length) {
-          createSessionAndRedirect(values.dataSourceIds);
-        }
-      })
-      .catch(() => null);
-  };
-
   if (activeSession) {
     return null;
   }
 
   if (dataSourcesStatus === "success" && dataSources.length > 0) {
-    return (
-      <PlaceholderContainer message="Start chatting with an existing Knowledge Base">
-        <Form autoCorrect="off" form={form} clearOnDestroy={true}>
-          <Flex gap={8}>
-            <Form.Item
-              name="dataSourceIds"
-              rules={[
-                { required: true, message: "Please select Knowledge Base(s)" },
-              ]}
-            >
-              <Select
-                mode="multiple"
-                disabled={dataSources.length === 0}
-                style={{ width: 300 }}
-                options={dataSources.map((value) => {
-                  return formatDataSource(value);
-                })}
-              />
-            </Form.Item>
-            <Button
-              type="primary"
-              icon={<ArrowRightOutlined />}
-              onClick={handleCreateSession}
-            />
-          </Flex>
-        </Form>
-      </PlaceholderContainer>
-    );
+    return null;
   }
 
   return (
-    <PlaceholderContainer message="Or create a knowledge base to chat with your documents.">
+    <PlaceholderContainer message="Create a knowledge base to chat with your documents.">
       <Button
         type="default"
         style={{ width: 200 }}

ui/src/pages/RagChatTab/ChatOutput/hooks/useCreateSessionAndRedirect.tsx

@@ -32,17 +32,24 @@ const useCreateSessionAndRedirect = (
     },
   });
 
-  return (dataSourceIds: number[], question?: string) => {
+  return (
+    dataSourceIds: number[],
+    question?: string,
+    inferenceModel?: string,
+  ) => {
     if (models) {
+      const supportsToolCalling = models.find(
+        (model) => model.model_id === inferenceModel,
+      )?.tool_calling_supported;
       const requestBody: CreateSessionRequest = {
         name: "",
         dataSourceIds: dataSourceIds,
-        inferenceModel: models[0].model_id,
+        inferenceModel: inferenceModel ?? models[0].model_id,
         responseChunks: 10,
         queryConfiguration: {
           enableHyde: false,
           enableSummaryFilter: true,
-          enableToolCalling: false,
+          enableToolCalling: supportsToolCalling ?? false,
           selectedTools: [],
         },
         embeddingModel: embeddingModels?.length

ui/src/pages/RagChatTab/FooterComponents/RagChatQueryInput.test.tsx

@@ -58,6 +58,26 @@ vi.mock("src/api/chatApi.ts", () => ({
   })),
 }));
 
+vi.mock("src/api/modelsApi.ts", () => ({
+  getLlmModelsQueryOptions: {
+    queryKey: ["llmModels"],
+    queryFn: () =>
+      Promise.resolve([
+        { model_id: "test-llm", tool_calling_supported: false },
+      ]),
+  },
+  useGetLlmModels: vi.fn(() => ({
+    data: [{ model_id: "test-llm" }],
+    isFetching: false,
+    error: null,
+  })),
+  useGetModelSource: vi.fn(() => ({
+    data: "OpenAI",
+    isFetching: false,
+    error: null,
+  })),
+}));
+
 vi.mock("src/api/ragQueryApi.ts", () => ({
   useSuggestQuestions: vi.fn(() => ({
     data: { suggested_questions: ["Sample question 1", "Sample question 2"] },
@@ -71,6 +91,16 @@ vi.mock("@tanstack/react-router", () => ({
   useSearch: vi.fn(() => ({ question: undefined })),
 }));
 
+vi.mock("@tanstack/react-query", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("@tanstack/react-query")>();
+  return {
+    ...actual,
+    useSuspenseQuery: vi.fn(() => ({
+      data: [{ model_id: "test-llm", tool_calling_supported: false }],
+    })),
+  };
+});
+
 vi.mock("src/utils/useModal.ts", () => ({
   default: vi.fn(() => ({
     isModalOpen: false,
@@ -472,22 +502,22 @@ describe("RagChatQueryInput", () => {
   });
 
   describe("New Session Callback", () => {
-    it("calls newSessionCallback when no sessionId exists", async () => {
+    it("calls newSessionCallback when no session exists", async () => {
       const user = userEvent.setup();
-      const mockNewSessionCallback = vi.fn();
+      const newSessionCallback = vi.fn();
       (useParams as Mock).mockReturnValue({ sessionId: undefined });
-
       const mockContext = createMockContext();
-      renderWithContext(mockContext, mockNewSessionCallback);
+      renderWithContext(mockContext, newSessionCallback);
 
       const textArea = screen.getByPlaceholderText("Ask a question");
-      const sendButton = screen.getByRole("button", { name: /send/i });
-
-      await user.type(textArea, "New session query");
-      await user.click(sendButton);
+      await user.type(textArea, "Test query");
+      await user.click(screen.getByRole("button", { name: /send/i }));
 
-      expect(mockNewSessionCallback).toHaveBeenCalledWith("New session query");
-      expect(mockStreamingChatMutation.mutate).not.toHaveBeenCalled();
+      expect(newSessionCallback).toHaveBeenCalledWith({
+        userInput: "Test query",
+        selectedDataSourceIds: [],
+        inferenceModel: "test-llm",
+      });
     });
   });