feat(kb): implement Document controller

internal/controller/document_controller.go

@@ -18,13 +18,27 @@ package controller
 
 import (
 	"context"
+	"fmt"
+	"strings"
 
+	batchv1 "k8s.io/api/batch/v1"
+	corev1 "k8s.io/api/core/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
+	apimeta "k8s.io/apimachinery/pkg/api/meta"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
 	logf "sigs.k8s.io/controller-runtime/pkg/log"
 
 	"github.com/LocoStack/loco-operator/api/v1alpha1"
+	"github.com/LocoStack/loco-operator/internal/reconciler"
+	"github.com/LocoStack/loco-operator/pkg/templates"
+)
+
+const (
+	finalizerDocumentCleanup = "locostack.com/document-cleanup"
 )
 
 // DocumentReconciler reconciles a Document object
@@ -36,28 +50,275 @@ type DocumentReconciler struct {
 // +kubebuilder:rbac:groups=locostack.com,resources=documents,verbs=get;list;watch;create;update;patch;delete
 // +kubebuilder:rbac:groups=locostack.com,resources=documents/status,verbs=get;update;patch
 // +kubebuilder:rbac:groups=locostack.com,resources=documents/finalizers,verbs=update
+// +kubebuilder:rbac:groups=locostack.com,resources=knowledgebases,verbs=get;list;watch
+// +kubebuilder:rbac:groups=locostack.com,resources=components,verbs=get;list;watch
+// +kubebuilder:rbac:groups=batch,resources=jobs,verbs=get;list;watch;create;update;patch;delete
 
-// Reconcile is part of the main kubernetes reconciliation loop which aims to
-// move the current state of the cluster closer to the desired state.
-// TODO(user): Modify the Reconcile function to compare the state specified by
-// the Document object against the actual cluster state, and then
-// perform operations to make the cluster state reflect the state specified by
-// the user.
-//
-// For more details, check Reconcile and its Result here:
-// - https://pkg.go.dev/sigs.k8s.io/controller-runtime@v0.23.3/pkg/reconcile
 func (r *DocumentReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
-	_ = logf.FromContext(ctx)
+	log := logf.FromContext(ctx)
+
+	doc := &v1alpha1.Document{}
+	if err := r.Get(ctx, req.NamespacedName, doc); err != nil {
+		if apierrors.IsNotFound(err) {
+			return ctrl.Result{}, nil
+		}
+		log.Error(err, "Failed to get Document", "namespace", req.NamespacedName.Namespace, "name", req.NamespacedName.Name)
+		return ctrl.Result{}, err
+	}
+
+	patch := client.MergeFrom(doc.DeepCopy())
+
+	ingestionJobTemplate, egestionJobTemplate, stack, endpoint, err := r.reconcileKB(ctx, doc)
+	if err != nil {
+		return ctrl.Result{}, err
+	}
+
+	if !controllerutil.ContainsFinalizer(doc, finalizerDocumentCleanup) {
+		patch := client.MergeFrom(doc.DeepCopy())
+		controllerutil.AddFinalizer(doc, finalizerDocumentCleanup)
+		return ctrl.Result{}, r.Patch(ctx, doc, patch)
+	}
 
-	// TODO(user): your logic here
+	if !doc.DeletionTimestamp.IsZero() {
+		egested, egestionFailed, err := r.reconcileEgestion(ctx, doc, egestionJobTemplate, stack, endpoint)
+		if err != nil {
+			return ctrl.Result{}, err
+		}
+		if !egested {
+			return ctrl.Result{}, nil
+		}
+		if egestionFailed {
+			// Handle egestion failure if necessary
+			return ctrl.Result{}, fmt.Errorf("Egestion failed for document %s/%s", req.NamespacedName.Namespace, req.NamespacedName.Name)
+		}
+		controllerutil.RemoveFinalizer(doc, finalizerDocumentCleanup)
+		return ctrl.Result{}, r.Patch(ctx, doc, patch)
+	}
 
+	ingested, ingestionFailed, err := r.reconcileIngestion(ctx, doc, ingestionJobTemplate, stack, endpoint)
+	if err != nil {
+		return ctrl.Result{}, err
+	}
+
+	if ingested {
+		apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
+			Type:               "Ready",
+			Status:             metav1.ConditionTrue,
+			Reason:             "Ingested",
+			Message:            "Document has been ingested into the KnowledgeBase",
+			ObservedGeneration: doc.Generation,
+		})
+	} else {
+		if ingestionFailed {
+			apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
+				Type:               "Ready",
+				Status:             metav1.ConditionFalse,
+				Reason:             "Failed",
+				Message:            "Document failed to be ingested into the KnowledgeBase",
+				ObservedGeneration: doc.Generation,
+			})
+		} else {
+			apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
+				Type:               "Ready",
+				Status:             metav1.ConditionFalse,
+				Reason:             "Ingesting",
+				Message:            "Document is being ingested into the KnowledgeBase",
+				ObservedGeneration: doc.Generation,
+			})
+		}
+
+	}
+
+	doc.Status.ObservedGeneration = doc.Generation
+	if err := r.Status().Patch(ctx, doc, patch); err != nil {
+		return ctrl.Result{}, client.IgnoreNotFound(err)
+	}
+
+	log.Info("Reconciled Document", "namespace", req.NamespacedName.Namespace, "name", req.NamespacedName.Name)
 	return ctrl.Result{}, nil
 }
 
+func (r *DocumentReconciler) reconcileKB(ctx context.Context, doc *v1alpha1.Document) (*v1alpha1.Template, *v1alpha1.Template, *v1alpha1.Stack, string, error) {
+	log := logf.FromContext(ctx)
+	kbRef := doc.Spec.KnowledgeBaseRef
+	if kbRef == nil {
+		// TODO
+		return nil, nil, nil, "", fmt.Errorf("KnowledgeBaseRef is required in Document spec")
+	}
+	var stackRef *corev1.LocalObjectReference
+	var ingestionJobTemplate *v1alpha1.Template
+	var egestionJobTemplate *v1alpha1.Template
+	var endpoint string
+	if kbRef.Kind == "ExternalKnowledgeBase" {
+		// TODO
+	} else if kbRef.Kind == "ManagedKnowledgeBase" {
+		kb := &v1alpha1.ManagedKnowledgeBase{}
+		if err := r.Get(ctx, client.ObjectKey{Namespace: doc.Namespace, Name: doc.Spec.KnowledgeBaseRef.Name}, kb); err != nil {
+			if apierrors.IsNotFound(err) {
+				apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
+					Type:               "KnowledgeBaseAvailable",
+					Status:             metav1.ConditionFalse,
+					Reason:             "KnowledgeBaseNotFound",
+					Message:            "Referenced KnowledgeBase not found",
+					ObservedGeneration: doc.Generation,
+				})
+				return nil, nil, nil, "", nil
+			}
+			log.Error(err, "Failed to get referenced ManagedKnowledgeBase", "namespace", doc.Namespace, "name", doc.Spec.KnowledgeBaseRef.Name)
+			return nil, nil, nil, "", err
+		}
+		stackRef = kb.Spec.StackRef
+		ingestionJobTemplate = kb.Spec.IngestionJob
+		egestionJobTemplate = kb.Spec.EgestionJob
+		endpoint = kb.Status.Endpoint
+	}
+	stack := &v1alpha1.Stack{}
+	if err := r.Get(ctx, client.ObjectKey{Namespace: doc.Namespace, Name: stackRef.Name}, stack); err != nil {
+		if apierrors.IsNotFound(err) {
+			apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
+				Type:               "KnowledgeBaseAvailable",
+				Status:             metav1.ConditionFalse,
+				Reason:             "StackNotFound",
+				Message:            "Stack for referenced KnowledgeBase not found",
+				ObservedGeneration: doc.Generation,
+			})
+			return nil, nil, nil, "", nil
+		}
+		log.Error(err, "Failed to get Stack for referenced ManagedKnowledgeBase", "namespace", doc.Namespace, "name", stackRef.Name)
+		return nil, nil, nil, "", err
+	}
+	apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
+		Type:               "KnowledgeBaseAvailable",
+		Status:             metav1.ConditionTrue,
+		Reason:             "KnowledgeBaseAvailable",
+		Message:            "Referenced KnowledgeBase is available",
+		ObservedGeneration: doc.Generation,
+	})
+	return ingestionJobTemplate, egestionJobTemplate, stack, endpoint, nil
+}
+
+func (r *DocumentReconciler) reconcileArtifacts(ctx context.Context, stack *v1alpha1.Stack, doc *v1alpha1.Document) (*corev1.PersistentVolumeClaim, bool, bool, error) {
+	artifactsReconciler := reconciler.NewArtifactReconciler(r.Client, r.Scheme, stack, "document", doc, &doc.Spec.Source)
+	return artifactsReconciler.ReconcileArtifact(ctx)
+}
+
+func (r *DocumentReconciler) reconcileIngestion(ctx context.Context, doc *v1alpha1.Document, jobTmpl *v1alpha1.Template, stack *v1alpha1.Stack, kbEndpoint string) (bool, bool, error) {
+	if doc.Status.Phase == "Ready" || doc.Status.Phase == "Failed" {
+		return true, doc.Status.Phase == "Failed", nil
+	}
+
+	tmpl, err := templates.Manager.ResolveTemplate(jobTmpl, "")
+	if err != nil {
+		apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
+			Type:               "IngestionJobResolved",
+			Status:             metav1.ConditionFalse,
+			Reason:             "TemplateResolutionFailed",
+			Message:            "Failed to resolve ingestion job template: " + err.Error(),
+			ObservedGeneration: doc.Generation,
+		})
+		return false, true, err
+	}
+	apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
+		Type:               "IngestionJobResolved",
+		Status:             metav1.ConditionTrue,
+		Reason:             "TemplateResolved",
+		Message:            "Ingestion job template resolved successfully",
+		ObservedGeneration: doc.Generation,
+	})
+
+	sourcePvc, sourceReady, sourceFailed, err := r.reconcileArtifacts(ctx, stack, doc)
+	if err != nil {
+		return false, true, err
+	}
+	if sourceFailed {
+		apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
+			Type:               "SourceReady",
+			Status:             metav1.ConditionFalse,
+			Reason:             "SourceFailed",
+			Message:            "Document source is not available",
+			ObservedGeneration: doc.Generation,
+		})
+		return false, true, nil
+	} else if !sourceReady {
+		apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
+			Type:               "SourceReady",
+			Status:             metav1.ConditionFalse,
+			Reason:             "SourceNotReady",
+			Message:            "Document source is not ready",
+			ObservedGeneration: doc.Generation,
+		})
+		return false, false, nil
+	}
+	apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
+		Type:               "SourceReady",
+		Status:             metav1.ConditionTrue,
+		Reason:             "SourceReady",
+		Message:            "Document source is ready",
+		ObservedGeneration: doc.Generation,
+	})
+
+	doc.Status.Phase = "Ingesting"
+	jobReconciler := reconciler.NewJobReconciler(r.Client, r.Scheme, stack, "DocumentIngestion", doc.Name, doc)
+	variables := map[string]string{
+		"spec.runtime.pvc":                   sourcePvc.Name,
+		"spec.docId":                         doc.Name,
+		"spec.chunkingStrategy.strategy":     doc.Spec.ChunkingStrategy.Strategy,
+		"spec.chunkingStrategy.chunkSize":    fmt.Sprint(doc.Spec.ChunkingStrategy.ChunkSize),
+		"spec.chunkingStrategy.chunkOverlap": fmt.Sprint(doc.Spec.ChunkingStrategy.ChunkOverlap),
+		"spec.chunkingStrategy.separators":   strings.Join(doc.Spec.ChunkingStrategy.Separators, ","),
+		"spec.kbEndpoint":                    kbEndpoint,
+	}
+	if doc.Spec.Source.PVC != nil {
+		variables["spec.filePath"] = doc.Spec.Source.PVC.FilePath
+	}
+	done, failed, err := jobReconciler.ReconcileJob(ctx, &tmpl, variables)
+	if done {
+		doc.Status.Phase = "Ready"
+	} else if failed {
+		doc.Status.Phase = "Failed"
+	}
+	return done, failed, err
+}
+
+func (r *DocumentReconciler) reconcileEgestion(ctx context.Context, doc *v1alpha1.Document, jobTmpl *v1alpha1.Template, stack *v1alpha1.Stack, kbEndpoint string) (bool, bool, error) {
+	if doc.Status.Phase != "Ready" {
+		// If the document is not in the "Ready" phase, we do not need egestion.
+		return true, false, nil
+	}
+
+	doc.Status.Phase = "Egesting"
+	jobReconciler := reconciler.NewJobReconciler(r.Client, r.Scheme, stack, "DocumentEgestion", doc.Name, doc)
+	tmpl, err := templates.Manager.ResolveTemplate(jobTmpl, "")
+	if err != nil {
+		apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
+			Type:               "EgestionJobResolved",
+			Status:             metav1.ConditionFalse,
+			Reason:             "TemplateResolutionFailed",
+			Message:            "Failed to resolve egestion job template: " + err.Error(),
+			ObservedGeneration: doc.Generation,
+		})
+		return false, true, err
+	}
+	apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
+		Type:               "EgestionJobResolved",
+		Status:             metav1.ConditionTrue,
+		Reason:             "TemplateResolved",
+		Message:            "Egestion job template resolved successfully",
+		ObservedGeneration: doc.Generation,
+	})
+	variables := map[string]string{
+		"spec.docId":      doc.Name,
+		"spec.kbEndpoint": kbEndpoint,
+	}
+	return jobReconciler.ReconcileJob(ctx, &tmpl, variables)
+}
+
 // SetupWithManager sets up the controller with the Manager.
 func (r *DocumentReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	return ctrl.NewControllerManagedBy(mgr).
 		For(&v1alpha1.Document{}).
+		Owns(&batchv1.Job{}).
+		Owns(&corev1.PersistentVolumeClaim{}).
 		Named("document").
 		Complete(r)
 }