From 74c738fb2ca0a3107f9ca1367de3a274f588bc0a Mon Sep 17 00:00:00 2001
From: Jussi Hallila <jussi@hallila.com>
Date: Thu, 3 Oct 2024 10:04:26 +0200
Subject: [PATCH] Add some more error handling to RAG UI Modal

---
 .changeset/good-vans-remain.md                |  5 ++
 plugins/frontend/rag-ai/src/api/client.ts     | 55 ++++++++++++-------
 .../src/components/RagModal/RagModal.tsx      | 10 +++-
 3 files changed, 50 insertions(+), 20 deletions(-)
 create mode 100644 .changeset/good-vans-remain.md

diff --git a/.changeset/good-vans-remain.md b/.changeset/good-vans-remain.md
new file mode 100644
index 000000000..fff84f801
--- /dev/null
+++ b/.changeset/good-vans-remain.md
@@ -0,0 +1,5 @@
+---
+'@roadiehq/rag-ai': patch
+---
+
+Add some error handling for better messages on RAG AI modal.
diff --git a/plugins/frontend/rag-ai/src/api/client.ts b/plugins/frontend/rag-ai/src/api/client.ts
index c14bdcae1..01b09d9d1 100644
--- a/plugins/frontend/rag-ai/src/api/client.ts
+++ b/plugins/frontend/rag-ai/src/api/client.ts
@@ -67,26 +67,43 @@ export class RoadieRagAiClient implements RagAiApi {
   async *ask(question: string, source: string): AsyncGenerator<ParsedEvent> {
     const { token } = await this.identityApi.getCredentials();
 
-    const stream = await this.fetch(`query/${source}`, {
-      body: JSON.stringify({
-        query: question,
-      }),
-      method: 'POST',
-      headers: {
-        'Content-Type': 'application/json',
-        Authorization: `Bearer ${token}`,
-      },
-    });
+    try {
+      const stream = await this.fetch(`query/${source}`, {
+        body: JSON.stringify({
+          query: question,
+        }),
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+          Authorization: `Bearer ${token}`,
+        },
+      });
+      if (stream) {
+        const reader = stream
+          .pipeThrough(new TextDecoderStream())
+          .pipeThrough(new EventSourceParserStream())
+          .getReader();
 
-    const reader = stream
-      .pipeThrough(new TextDecoderStream())
-      .pipeThrough(new EventSourceParserStream())
-      .getReader();
-
-    while (true) {
-      const { done, value } = await reader.read();
-      if (done) break;
-      yield value;
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          yield value;
+        }
+      } else {
+        yield {
+          data: 'No response received from the LLM',
+          event: 'message',
+          type: 'event',
+        };
+      }
+    } catch (e: any) {
+      // eslint-disable-next-line
+      console.error(e.message);
+      yield {
+        data: `Failed to receive response from the backend endpoint.`,
+        event: 'error',
+        type: 'event',
+      };
     }
   }
 }
diff --git a/plugins/frontend/rag-ai/src/components/RagModal/RagModal.tsx b/plugins/frontend/rag-ai/src/components/RagModal/RagModal.tsx
index 688c4112c..c9c5f4265 100644
--- a/plugins/frontend/rag-ai/src/components/RagModal/RagModal.tsx
+++ b/plugins/frontend/rag-ai/src/components/RagModal/RagModal.tsx
@@ -33,6 +33,7 @@ import { ragAiApiRef } from '../../api';
 import { useApi } from '@backstage/core-plugin-api';
 import { ResponseEmbedding } from '../../types';
 import { Thinking } from './Thinking';
+import { WarningPanel } from '@backstage/core-components';
 
 export type RagModalProps = {
   title?: string;
@@ -88,11 +89,13 @@ export const ControlledRagModal = ({
   const [thinking, setThinking] = useState(false);
   const [questionResult, setQuestionResult] = useState('');
   const [embeddings, setEmbeddings] = useState<ResponseEmbedding[]>([]);
+  const [warning, setWarning] = useState<string | undefined>();
   const ragApi = useApi(ragAiApiRef);
   const askLlm = useCallback(
     async (question: string, source: string) => {
       setThinking(true);
       setQuestionResult('');
+      setWarning(undefined);
       setEmbeddings([]);
 
       for await (const chunk of ragApi.ask(question, source)) {
@@ -105,6 +108,10 @@ export const ControlledRagModal = ({
             setEmbeddings(JSON.parse(chunk.data));
             break;
           }
+          case 'error': {
+            setWarning(chunk.data);
+            break;
+          }
           default:
             throw new Error(`Unknown event type: ${chunk.event}`);
         }
@@ -148,7 +155,8 @@ export const ControlledRagModal = ({
             }}
           />
         </Box>
-        {thinking && !questionResult ? (
+        {warning && <WarningPanel severity="warning" message={warning} />}
+        {thinking && !questionResult && !warning ? (
           <Box p={6} display="flex" justifyContent="center" alignItems="center">
             <Thinking />
           </Box>