Chore: use SSE to deliver change event for knowledgefiles (#1106)

* Chore: use SSE to deliver change event for knowledgefiles

Signed-off-by: Daishan Peng <[email protected]>

pkg/api/handlers/agent.go

@@ -2,6 +2,7 @@ package handlers
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"net/http"
@@ -15,9 +16,11 @@ import (
 	"github.com/obot-platform/obot/pkg/invoke"
 	"github.com/obot-platform/obot/pkg/render"
 	v1 "github.com/obot-platform/obot/pkg/storage/apis/otto.otto8.ai/v1"
+	"github.com/obot-platform/obot/pkg/storage/selectors"
 	"github.com/obot-platform/obot/pkg/system"
 	"github.com/obot-platform/obot/pkg/wait"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/fields"
 	kclient "sigs.k8s.io/controller-runtime/pkg/client"
 )
 
@@ -844,6 +847,71 @@ func (a *AgentHandler) Script(req api.Context) error {
 	return req.Write(script)
 }
 
+func (a *AgentHandler) WatchKnowledgeFile(req api.Context) error {
+	knowledgeSetNames, agentName, err := a.getKnowledgeSetsAndName(req, req.PathValue("agent_id"))
+	if err != nil {
+		return err
+	}
+
+	if len(knowledgeSetNames) == 0 {
+		return req.Write(types.KnowledgeFileList{Items: []types.KnowledgeFile{}})
+	}
+
+	knowledgeSourceName := req.PathValue("knowledge_source_id")
+	var knowledgeSource *v1.KnowledgeSource
+	if knowledgeSourceName != "" {
+		knowledgeSource = &v1.KnowledgeSource{}
+		if err := req.Get(knowledgeSource, knowledgeSourceName); err != nil {
+			return err
+		}
+		if knowledgeSource.Spec.KnowledgeSetName != knowledgeSetNames[0] {
+			return types.NewErrBadRequest("knowledgeSource %q does not belong to agent %q", knowledgeSource.Name, agentName)
+		}
+	}
+
+	w, err := req.Storage.Watch(req.Context(), &v1.KnowledgeFileList{}, kclient.InNamespace(req.Namespace()),
+		&kclient.ListOptions{
+			FieldSelector: fields.SelectorFromSet(selectors.RemoveEmpty(map[string]string{
+				"spec.knowledgeSetName":    knowledgeSetNames[0],
+				"spec.knowledgeSourceName": knowledgeSourceName,
+			})),
+		})
+	if err != nil {
+		return err
+	}
+	defer func() {
+		w.Stop()
+		//nolint:revive
+		for range w.ResultChan() {
+		}
+	}()
+
+	req.ResponseWriter.Header().Set("Content-Type", "text/event-stream")
+	defer func() {
+		_ = req.WriteDataEvent(api.EventClose{})
+	}()
+
+	for event := range w.ResultChan() {
+		if knowledgeFile, ok := event.Object.(*v1.KnowledgeFile); ok {
+			payload := map[string]any{
+				"eventType":     event.Type,
+				"knowledgeFile": convertKnowledgeFile(agentName, "", *knowledgeFile),
+			}
+			data, err := json.Marshal(payload)
+			if err != nil {
+				return err
+			}
+			sseEvent := fmt.Sprintf("data: %s\n\n", data)
+			if _, err := req.ResponseWriter.Write([]byte(sseEvent)); err != nil {
+				return err
+			}
+			req.Flush()
+		}
+	}
+
+	return nil
+}
+
 func MetadataFrom(obj kclient.Object, linkKV ...string) types.Metadata {
 	m := types.Metadata{
 		ID:       obj.GetName(),

pkg/api/router/router.go

@@ -141,6 +141,7 @@ func Router(services *services.Services) (http.Handler, error) {
 	mux.HandleFunc("DELETE /api/agents/{agent_id}/knowledge-sources/{id}", agents.DeleteKnowledgeSource)
 	mux.HandleFunc("GET /api/agents/{agent_id}/knowledge-sources/{knowledge_source_id}/knowledge-files", agents.ListKnowledgeFiles)
 	mux.HandleFunc("POST /api/agents/{agent_id}/knowledge-sources/{knowledge_source_id}/knowledge-files/{file_id}/ingest", agents.ReIngestKnowledgeFile)
+	mux.HandleFunc("GET /api/agents/{agent_id}/knowledge-sources/{knowledge_source_id}/knowledge-files/watch", agents.WatchKnowledgeFile)
 
 	// Workflows
 	mux.HandleFunc("GET /api/workflows", workflows.List)

ui/admin/app/components/knowledge/FileTree.tsx

@@ -141,9 +141,6 @@ export default function FileTreeNode({
             filePathPrefixInclude,
             filePathPrefixExclude,
         });
-
-        // We should manually approve/unapprove all files in the folder at once so that we don't rely on backend reconciliation logic as it will cause delay in updating the UI.
-        onApproveFile(!included, node);
     };
 
     return (

ui/admin/app/hooks/knowledge/useKnowledgeSourceFiles.ts

@@ -1,8 +1,8 @@
 import { useEffect, useMemo, useState } from "react";
-import useSWR from "swr";
 
 import {
     KnowledgeFile,
+    KnowledgeFileEvent,
     KnowledgeFileState,
     KnowledgeSource,
     KnowledgeSourceNamespace,
@@ -29,27 +29,56 @@ export function useKnowledgeSourceFiles(
         startPolling();
     }
 
-    const {
-        data: files,
-        mutate: mutateFiles,
-        ...rest
-    } = useSWR(
-        KnowledgeSourceApiService.getFilesForKnowledgeSource.key(
-            namespace,
-            agentId,
-            knowledgeSource.id
-        ),
-        ({ agentId, sourceId }) =>
-            KnowledgeSourceApiService.getFilesForKnowledgeSource(
+    const [files, setFiles] = useState<KnowledgeFile[]>([]);
+
+    useEffect(() => {
+        const eventSource =
+            KnowledgeSourceApiService.getKnowledgeSourceFilesEventSource(
                 namespace,
                 agentId,
-                sourceId
-            ),
-        {
-            revalidateOnFocus: false,
-            refreshInterval: blockPollingFiles ? undefined : 5000,
-        }
-    );
+                knowledgeSource.id
+            );
+
+        eventSource.onmessage = (event) => {
+            const { eventType, knowledgeFile } = JSON.parse(
+                event.data
+            ) as KnowledgeFileEvent;
+
+            setFiles((prevFiles) => {
+                let updatedFiles = [...prevFiles];
+                switch (eventType) {
+                    case "ADDED":
+                    case "MODIFIED":
+                        {
+                            const existingIndex = updatedFiles.findIndex(
+                                (file) => file.id === knowledgeFile.id
+                            );
+                            if (existingIndex !== -1) {
+                                updatedFiles[existingIndex] = knowledgeFile;
+                            } else {
+                                updatedFiles.push(knowledgeFile);
+                            }
+                        }
+                        break;
+                    case "DELETED":
+                        {
+                            updatedFiles = updatedFiles.filter(
+                                (file) => file.id !== knowledgeFile.id
+                            );
+                        }
+                        break;
+                    default:
+                        break;
+                }
+                return updatedFiles;
+            });
+        };
+
+        return () => {
+            setFiles([]);
+            eventSource.close();
+        };
+    }, [knowledgeSource.id, namespace, agentId]);
 
     const sortedFiles = useMemo(() => {
         return (
@@ -90,7 +119,7 @@ export function useKnowledgeSourceFiles(
                 knowledgeSource.id,
                 fileId
             );
-        mutateFiles((prev) =>
+        setFiles((prev) =>
             prev?.map((f) => (f.id === fileId ? updatedFile : f))
         );
     };
@@ -109,7 +138,7 @@ export function useKnowledgeSourceFiles(
             console.error("Failed to approve file", error);
         }
 
-        mutateFiles((prev) =>
+        setFiles((prev) =>
             prev?.map((f) => (f.id === file.id ? (updatedFile ?? file) : f))
         );
     };
@@ -118,8 +147,6 @@ export function useKnowledgeSourceFiles(
         files: sortedFiles,
         reingestFile,
         approveFile,
-        mutateFiles,
         startPollingFiles: startPolling,
-        ...rest,
     };
 }

ui/admin/app/lib/model/knowledge.ts

@@ -153,6 +153,11 @@ export type KnowledgeFile = {
     sizeInBytes?: number;
 };
 
+export type KnowledgeFileEvent = {
+    eventType: "ADDED" | "MODIFIED" | "DELETED";
+    knowledgeFile: KnowledgeFile;
+};
+
 export function getRemoteFileDisplayName(item: KnowledgeFile) {
     return item.fileName;
 }

ui/admin/app/lib/routers/apiRoutes.ts

@@ -122,6 +122,14 @@ export const ApiRoutes = {
             entityId: string,
             fileName: string
         ) => buildUrl(`/${namespace}/${entityId}/approve-file/${fileName}`),
+        watchKnowledgeSourceFiles: (
+            namespace: KnowledgeSourceNamespace,
+            entityId: string,
+            sourceId: string
+        ) =>
+            buildUrl(
+                `/${namespace}/${entityId}/knowledge-sources/${sourceId}/knowledge-files/watch`
+            ),
     },
     knowledgeFiles: {
         getKnowledgeFiles: (

ui/admin/app/lib/service/api/knowledgeSourceApiService.ts

@@ -176,13 +176,28 @@ async function deleteKnowledgeSource(
     });
 }
 
+function getKnowledgeSourceFilesEventSource(
+    namespace: KnowledgeSourceNamespace,
+    agentId: string,
+    sourceId: string
+) {
+    return new EventSource(
+        ApiRoutes.knowledgeSources.watchKnowledgeSourceFiles(
+            namespace,
+            agentId,
+            sourceId
+        ).url
+    );
+}
+
 export const KnowledgeSourceApiService = {
     approveFile,
     createKnowledgeSource,
     updateKnowledgeSource,
     resyncKnowledgeSource,
     getKnowledgeSources,
     getFilesForKnowledgeSource,
+    getKnowledgeSourceFilesEventSource,
     reingestFileFromSource,
     deleteKnowledgeSource,
 };