Feat: Add file tree, add include and exclude prefix for auto ingestion

Signed-off-by: Daishan Peng <[email protected]>
otto8-ai · Nov 20, 2024 · 58ed35e · 58ed35e
1 parent fd072b4
commit 58ed35e
Show file tree

Hide file tree

Showing 9 changed files with 681 additions and 590 deletions.
diff --git a/apiclient/types/knowledgesource.go b/apiclient/types/knowledgesource.go
@@ -36,12 +36,16 @@ type KnowledgeSource struct {
 	LastSyncStartTime       *Time                `json:"lastSyncStartTime,omitempty"`
 	LastSyncEndTime         *Time                `json:"lastSyncEndTime,omitempty"`
 	LastRunID               string               `json:"lastRunID,omitempty"`
+	FilePathPrefixInclude   []string             `json:"filePathPrefixInclude,omitempty"`
+	FilePathPrefixExclude   []string             `json:"filePathPrefixExclude,omitempty"`
 }
 
 type KnowledgeSourceManifest struct {
-	SyncSchedule         string `json:"syncSchedule,omitempty"`
-	AutoApprove          *bool  `json:"autoApprove,omitempty"`
-	KnowledgeSourceInput `json:",inline"`
+	SyncSchedule          string   `json:"syncSchedule,omitempty"`
+	AutoApprove           *bool    `json:"autoApprove,omitempty"`
+	FilePathPrefixInclude []string `json:"filePathPrefixInclude,omitempty"`
+	FilePathPrefixExclude []string `json:"filePathPrefixExclude,omitempty"`
+	KnowledgeSourceInput  `json:",inline"`
 }
 
 type KnowledgeSourceList List[KnowledgeSource]

diff --git a/pkg/api/handlers/knowledgesource.go b/pkg/api/handlers/knowledgesource.go
@@ -25,6 +25,8 @@ func convertKnowledgeSource(agentName string, knowledgeSource v1.KnowledgeSource
 		LastSyncStartTime:       types.NewTime(knowledgeSource.Status.LastSyncStartTime.Time),
 		LastSyncEndTime:         types.NewTime(knowledgeSource.Status.LastSyncEndTime.Time),
 		LastRunID:               knowledgeSource.Status.RunName,
+		FilePathPrefixExclude:   knowledgeSource.Spec.Manifest.FilePathPrefixExclude,
+		FilePathPrefixInclude:   knowledgeSource.Spec.Manifest.FilePathPrefixInclude,
 	}
 }
 

diff --git a/pkg/controller/handlers/knowledgefile/knowledgefile.go b/pkg/controller/handlers/knowledgefile/knowledgefile.go
@@ -99,6 +99,30 @@ func (h *Handler) IngestFile(req router.Request, _ router.Response) error {
 		}
 	}
 
+	// Check approval
+	matchInclude := isFileMatchIncludePattern(file.Spec.FileName, source.Spec.Manifest.FilePathPrefixInclude)
+	matchExclude := isFileMatchExcludePattern(file.Spec.FileName, source.Spec.Manifest.FilePathPrefixExclude)
+	if file.Spec.Approved == nil {
+		if source.Spec.Manifest.AutoApprove != nil && *source.Spec.Manifest.AutoApprove {
+			file.Spec.Approved = typed.Pointer(true)
+			if err := req.Client.Update(req.Ctx, file); err != nil {
+				return err
+			}
+		}
+
+		if matchInclude && !matchExclude {
+			file.Spec.Approved = typed.Pointer(true)
+			if err := req.Client.Update(req.Ctx, file); err != nil {
+				return err
+			}
+		} else if matchExclude {
+			file.Spec.Approved = typed.Pointer(false)
+			if err := req.Client.Update(req.Ctx, file); err != nil {
+				return err
+			}
+		}
+	}
+
 	if file.Status.State.IsTerminal() && !shouldReIngest(file) {
 		return nil
 	}
@@ -111,16 +135,6 @@ func (h *Handler) IngestFile(req router.Request, _ router.Response) error {
 		}
 	}
 
-	// Check approval
-	if file.Spec.Approved == nil {
-		if source.Spec.Manifest.AutoApprove != nil && *source.Spec.Manifest.AutoApprove {
-			file.Spec.Approved = typed.Pointer(true)
-			if err := req.Client.Update(req.Ctx, file); err != nil {
-				return err
-			}
-		}
-	}
-
 	if file.Spec.Approved == nil || !*file.Spec.Approved {
 		// Not approved, wait for user action
 		return nil
@@ -402,3 +416,23 @@ func (h *Handler) Cleanup(req router.Request, _ router.Response) error {
 	}
 	return nil
 }
+
+func isFileMatchIncludePattern(filePath string, filePathPrefixInclude []string) bool {
+	for _, include := range filePathPrefixInclude {
+		if strings.HasPrefix(filePath, include) {
+			return true
+		}
+	}
+
+	return false
+}
+
+func isFileMatchExcludePattern(filePath string, filePathPrefixExclude []string) bool {
+	for _, exclude := range filePathPrefixExclude {
+		if strings.HasPrefix(filePath, exclude) {
+			return true
+		}
+	}
+
+	return false
+}