remove the datasource id from the context and use the active session …

…instead

ui/src/pages/RagChatTab/ChatLayout.tsx

@@ -53,25 +53,21 @@ const getSessionForSessionId = (sessionId?: string, sessions?: Session[]) => {
   return sessions?.find((session) => session.id.toString() === sessionId);
 };
 
-const getDataSourceIdForSession = (session?: Session) => {
-  return session?.dataSourceIds[0];
-};
-
 function ChatLayout() {
   const { data: sessions } = useSuspenseQuery(getSessionsQueryOptions);
 
   const { sessionId } = useParams({ strict: false });
-  const activeSession = getSessionForSessionId(sessionId, sessions);
-  const dataSourceId = getDataSourceIdForSession(activeSession);
   const [currentQuestion, setCurrentQuestion] = useState("");
-
   const { data: dataSources, status: dataSourcesStatus } =
     useGetDataSourcesQuery();
   const [excludeKnowledgeBase, setExcludeKnowledgeBase] = useState(false);
   const { status: chatHistoryStatus, data: chatHistory } = useChatHistoryQuery(
     sessionId?.toString() ?? "",
   );
 
+  const activeSession = getSessionForSessionId(sessionId, sessions);
+  const dataSourceId = activeSession?.dataSourceIds[0];
+
   const dataSourceSize = useMemo(() => {
     return (
       dataSources?.find((ds) => ds.id === dataSourceId)?.totalDocSize ?? null
@@ -90,7 +86,6 @@ function ChatLayout() {
   return (
     <RagChatContext.Provider
       value={{
-        dataSourceId,
         excludeKnowledgeBaseState: [
           excludeKnowledgeBase,
           setExcludeKnowledgeBase,

ui/src/pages/RagChatTab/ChatOutput/ChatMessages/ChatBodyController.test.tsx

@@ -96,7 +96,6 @@ describe("ChatBodyController", () => {
     const defaultContextValue: RagChatContextType = {
       chatHistoryQuery: { chatHistoryStatus: undefined, chatHistory: [] },
       currentQuestionState: ["", () => null],
-      dataSourceId: undefined,
       dataSourcesQuery: { dataSourcesStatus: undefined, dataSources: [] },
       excludeKnowledgeBaseState: [false, () => null],
       dataSourceSize: null,
@@ -125,7 +124,6 @@ describe("ChatBodyController", () => {
 
   it("renders NoSessionState when no sessionId and dataSources are available", () => {
     renderWithContext({
-      dataSourceId: undefined,
       dataSourcesQuery: { dataSourcesStatus: undefined, dataSources: [] },
       dataSourceSize: null,
       activeSession: undefined,
@@ -188,7 +186,6 @@ describe("ChatBodyController", () => {
   it("renders NoDataSourceForSession when no currentDataSource is found", () => {
     renderWithContext({
       dataSourcesQuery: { dataSources: [testDataSource] },
-      dataSourceId: undefined,
       activeSession: { ...testSession, dataSourceIds: [2] },
     });
 

ui/src/pages/RagChatTab/ChatOutput/ChatMessages/ChatMessageController.tsx

@@ -43,9 +43,10 @@ import { RagChatContext } from "pages/RagChatTab/State/RagChatContext.tsx";
 const ChatMessageController = () => {
   const {
     chatHistoryQuery: { chatHistory },
-    dataSourceId,
+    activeSession,
   } = useContext(RagChatContext);
   const scrollEl = useRef<HTMLDivElement>(null);
+  const dataSourceId = activeSession?.dataSourceIds[0];
 
   useEffect(() => {
     setTimeout(() => {

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

@@ -46,8 +46,8 @@ import { cdlGray700 } from "src/cuix/variables.ts";
 import Images from "src/components/images/Images.ts";
 
 const DataSourceSummaryCard = () => {
-  const { dataSourceId } = useContext(RagChatContext);
-
+  const { activeSession } = useContext(RagChatContext);
+  const dataSourceId = activeSession?.dataSourceIds[0];
   const dataSourceSummary = useGetDataSourceSummary({
     data_source_id: dataSourceId?.toString() ?? "",
     queryEnabled: true,
@@ -81,7 +81,9 @@ const DataSourceSummaryCard = () => {
 };
 
 const EmptyChatState = ({ dataSourceSize }: { dataSourceSize: number }) => {
-  const { dataSourceId } = useContext(RagChatContext);
+  const { activeSession } = useContext(RagChatContext);
+  const dataSourceId = activeSession?.dataSourceIds[0];
+
   return (
     <Flex vertical align="center" style={{ height: "100%" }}>
       <Image

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

@@ -45,12 +45,11 @@ import { createQueryConfiguration, useChatMutation } from "src/api/chatApi.ts";
 
 const SuggestedQuestionsCards = () => {
   const {
-    dataSourceId,
     currentQuestionState: [, setCurrentQuestion],
     activeSession,
     excludeKnowledgeBaseState: [excludeKnowledgeBase],
   } = useContext(RagChatContext);
-
+  const dataSourceId = activeSession?.dataSourceIds[0];
   const sessionId = activeSession?.id.toString();
   const {
     data,

ui/src/pages/RagChatTab/ChatOutput/Sources/SourceCard.tsx

@@ -57,9 +57,10 @@ import { cdlGray600 } from "src/cuix/variables.ts";
 import MetaData from "pages/RagChatTab/ChatOutput/Sources/MetaData.tsx";
 
 export const SourceCard = ({ source }: { source: SourceNode }) => {
-  const { dataSourceId } = useContext(RagChatContext);
+  const { activeSession } = useContext(RagChatContext);
   const [showContent, setShowContent] = useState(false);
   const chunkContents = useGetChunkContents();
+  const dataSourceId = activeSession?.dataSourceIds[0];
   const documentSummary = useGetDocumentSummary({
     data_source_id: dataSourceId?.toString() ?? "",
     doc_id: source.doc_id,

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

@@ -51,15 +51,14 @@ import type { SwitchChangeEventHandler } from "antd/lib/switch";
 
 const RagChatQueryInput = () => {
   const {
-    dataSourceId,
     excludeKnowledgeBaseState: [excludeKnowledgeBase, setExcludeKnowledgeBase],
     currentQuestionState: [, setCurrentQuestion],
     chatHistoryQuery: { chatHistory },
     dataSourceSize,
     dataSourcesQuery: { dataSourcesStatus },
     activeSession,
   } = useContext(RagChatContext);
-
+  const dataSourceId = activeSession?.dataSourceIds[0];
   const [userInput, setUserInput] = useState("");
   const { sessionId } = useParams({ strict: false });
 

ui/src/pages/RagChatTab/RagChat.tsx

@@ -47,13 +47,12 @@ const { Footer, Content } = Layout;
 
 const RagChat = () => {
   const {
-    dataSourceId,
     dataSourcesQuery: { dataSources },
     activeSession,
   } = useContext(RagChatContext);
 
   const currentDataSource = dataSources.find((dataSource) => {
-    return dataSource.id === dataSourceId;
+    return dataSource.id === activeSession?.dataSourceIds[0];
   });
 
   return (

ui/src/pages/RagChatTab/State/RagChatContext.tsx

@@ -52,7 +52,6 @@ export interface RagChatContextType {
     dataSources: DataSourceType[];
     dataSourcesStatus?: "error" | "success" | "pending";
   };
-  dataSourceId?: number;
   dataSourceSize: number | null;
   excludeKnowledgeBaseState: [boolean, Dispatch<SetStateAction<boolean>>];
 }
@@ -62,7 +61,6 @@ export const RagChatContext = createContext<RagChatContextType>({
   currentQuestionState: ["", () => null],
   chatHistoryQuery: { chatHistory: [], chatHistoryStatus: undefined },
   dataSourcesQuery: { dataSources: [], dataSourcesStatus: undefined },
-  dataSourceId: undefined, // TODO: remove this and have it pulled from active session
   dataSourceSize: null,
   excludeKnowledgeBaseState: [false, () => null],
 });