Display row number in the UI if present (#46)

* Display row number in the UI if present * lint
cloudera · Nov 22, 2024 · 50bdcd0 · 50bdcd0
1 parent ebbdf89
commit 50bdcd0
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 22 deletions.
diff --git a/llm-service/app/routers/index/data_source/__init__.py b/llm-service/app/routers/index/data_source/__init__.py
@@ -30,6 +30,7 @@
 
 import logging
 import tempfile
+from typing import Any, Dict
 
 from fastapi import APIRouter, Depends
 from fastapi_utils.cbv import cbv
@@ -65,6 +66,11 @@ class RagIndexDocumentRequest(BaseModel):
     configuration: RagIndexDocumentConfiguration = RagIndexDocumentConfiguration()
 
 
+class ChunkContentsResponse(BaseModel):
+    text: str
+    metadata: Dict[str, Any]
+
+
 @cbv(router)
 class DataSourceController:
     chunks_vector_store: VectorStore = Depends(
@@ -86,11 +92,11 @@ def size(self) -> int:
         response_model=None,
     )
     @exceptions.propagates
-    def chunk_contents(self, chunk_id: str) -> str:
-        return (
-            self.chunks_vector_store.llama_vector_store()
-            .get_nodes([chunk_id])[0]
-            .get_content()
+    def chunk_contents(self, chunk_id: str) -> ChunkContentsResponse:
+        node = self.chunks_vector_store.llama_vector_store().get_nodes([chunk_id])[0]
+        return ChunkContentsResponse(
+            text=node.get_content(),
+            metadata=node.metadata,
         )
 
     @router.delete(

diff --git a/ui/src/api/ragQueryApi.ts b/ui/src/api/ragQueryApi.ts
@@ -37,14 +37,14 @@
  ******************************************************************************/
 
 import { useMutation, useQuery } from "@tanstack/react-query";
+import { QueryConfiguration } from "src/api/chatApi.ts";
 import {
   getRequest,
   llmServicePath,
   MutationKeys,
   postRequest,
   QueryKeys,
 } from "src/api/utils.ts";
-import { QueryConfiguration } from "src/api/chatApi.ts";
 
 export interface RagMessage {
   role: "user" | "assistant";
@@ -94,7 +94,10 @@ const suggestQuestionsQuery = async (
   );
 };
 
-type ChunkContents = string;
+interface ChunkContentsResponse {
+  text: string;
+  metadata: Record<string, string | number>;
+}
 
 interface ChunkContentsRequest {
   data_source_id: string;
@@ -110,7 +113,7 @@ export const useGetChunkContents = () => {
 
 const getChunkContents = async (
   request: ChunkContentsRequest,
-): Promise<ChunkContents> => {
+): Promise<ChunkContentsResponse> => {
   return getRequest(
     `${llmServicePath}/data_sources/${request.data_source_id}/chunks/${request.chunk_id}`,
   );

diff --git a/ui/src/pages/RagChatTab/ChatOutput/Sources/SourceCard.tsx b/ui/src/pages/RagChatTab/ChatOutput/Sources/SourceCard.tsx
@@ -36,9 +36,7 @@
  * DATA.
  ******************************************************************************/
 
-import { useGetChunkContents } from "src/api/ragQueryApi.ts";
-import { useContext, useState } from "react";
-import { RagChatContext } from "pages/RagChatTab/State/RagChatContext.tsx";
+import Icon from "@ant-design/icons";
 import {
   Alert,
   Card,
@@ -49,10 +47,12 @@ import {
   Tooltip,
   Typography,
 } from "antd";
+import { RagChatContext } from "pages/RagChatTab/State/RagChatContext.tsx";
+import { useContext, useState } from "react";
 import { SourceNode } from "src/api/chatApi.ts";
+import { useGetChunkContents } from "src/api/ragQueryApi.ts";
 import { useGetDocumentSummary } from "src/api/summaryApi.ts";
 import DocumentationIcon from "src/cuix/icons/DocumentationIcon";
-import Icon from "@ant-design/icons";
 import { cdlGray600 } from "src/cuix/variables.ts";
 
 export const SourceCard = ({ source }: { source: SourceNode }) => {
@@ -122,16 +122,28 @@ export const SourceCard = ({ source }: { source: SourceNode }) => {
                 </div>
               </Flex>
             ) : (
-              <Flex vertical>
-                <Typography.Title level={5} style={{ marginTop: 10 }}>
-                  Extracted reference content
-                </Typography.Title>
-                <Typography.Paragraph
-                  style={{ textAlign: "left", whiteSpace: "pre-wrap" }}
-                >
-                  {chunkContents.data}
-                </Typography.Paragraph>
-              </Flex>
+              chunkContents.data && (
+                <Flex vertical>
+                  <Typography.Title level={5} style={{ marginTop: 10 }}>
+                    Extracted reference content
+                  </Typography.Title>
+                  <Typography.Paragraph
+                    style={{ textAlign: "left", whiteSpace: "pre-wrap" }}
+                  >
+                    {chunkContents.data.text}
+                  </Typography.Paragraph>
+                  {chunkContents.data.metadata.row_number && (
+                    <>
+                      <Typography.Title level={5} style={{ marginTop: 0 }}>
+                        Metadata
+                      </Typography.Title>
+                      <Typography.Text>
+                        Row number: {chunkContents.data.metadata.row_number}
+                      </Typography.Text>
+                    </>
+                  )}
+                </Flex>
+              )
             )}
           </Flex>
         </Card>