Prechádzať zdrojové kódy

feat(kb): implement ManagedKnowledgeBase controller

Thomas Zhang 2 mesiacov pred
rodič
commit
0e791b75db

+ 150 - 11
internal/controller/managedknowledgebase_controller.go

@@ -18,13 +18,25 @@ package controller
 
 import (
 	"context"
+	"fmt"
+	"maps"
 
+	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"
+	"k8s.io/apimachinery/pkg/types"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
+	"sigs.k8s.io/controller-runtime/pkg/handler"
 	logf "sigs.k8s.io/controller-runtime/pkg/log"
+	"sigs.k8s.io/controller-runtime/pkg/reconcile"
 
 	"github.com/LocoStack/loco-operator/api/v1alpha1"
+	"github.com/LocoStack/loco-operator/internal/reconciler"
+	"github.com/LocoStack/loco-operator/pkg/templates"
 )
 
 // ManagedKnowledgeBaseReconciler reconciles a ManagedKnowledgeBase object
@@ -36,28 +48,155 @@ type ManagedKnowledgeBaseReconciler struct {
 // +kubebuilder:rbac:groups=locostack.com,resources=managedknowledgebases,verbs=get;list;watch;create;update;patch;delete
 // +kubebuilder:rbac:groups=locostack.com,resources=managedknowledgebases/status,verbs=get;update;patch
 // +kubebuilder:rbac:groups=locostack.com,resources=managedknowledgebases/finalizers,verbs=update
+// +kubebuilder:rbac:groups=locostack.com,resources=components,verbs=get;list;watch;create;update;patch;delete
+// +kubebuilder:rbac:groups=locostack.com,resources=stacks,verbs=get;list;watch
 
-// 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 ManagedKnowledgeBase 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 *ManagedKnowledgeBaseReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
-	_ = logf.FromContext(ctx)
+	log := logf.FromContext(ctx)
 
-	// TODO(user): your logic here
+	mkb := &v1alpha1.ManagedKnowledgeBase{}
+	if err := r.Get(ctx, req.NamespacedName, mkb); err != nil {
+		if apierrors.IsNotFound(err) {
+			return ctrl.Result{}, nil
+		}
+		log.Error(err, "Failed to get ManagedKnowledgeBase", "namespace", req.NamespacedName.Namespace, "name", req.NamespacedName.Name)
+		return ctrl.Result{}, err
+	}
 
+	patch := client.MergeFrom(mkb.DeepCopy())
+
+	componentAvailable, err := r.reconcileComponent(ctx, mkb)
+	if err != nil {
+		return ctrl.Result{}, err
+	}
+	if !componentAvailable {
+		apimeta.SetStatusCondition(&mkb.Status.Conditions, metav1.Condition{
+			Type:               "Available",
+			Status:             metav1.ConditionFalse,
+			Reason:             "KnowledgeBaseUnavailable",
+			Message:            "ManagedKnowledgeBase is not available",
+			ObservedGeneration: mkb.Generation,
+		})
+	} else {
+		apimeta.SetStatusCondition(&mkb.Status.Conditions, metav1.Condition{
+			Type:               "Available",
+			Status:             metav1.ConditionTrue,
+			Reason:             "KnowledgeBaseAvailable",
+			Message:            "ManagedKnowledgeBase is available",
+			ObservedGeneration: mkb.Generation,
+		})
+	}
+
+	mkb.Status.ObservedGeneration = mkb.Generation
+	if err := r.Status().Patch(ctx, mkb, patch); err != nil {
+		return ctrl.Result{}, client.IgnoreNotFound(err)
+	}
+
+	log.Info("Reconciled ManagedKnowledgeBase", "namespace", req.NamespacedName.Namespace, "name", req.NamespacedName.Name)
 	return ctrl.Result{}, nil
 }
 
+func (r *ManagedKnowledgeBaseReconciler) reconcileComponent(ctx context.Context, mkb *v1alpha1.ManagedKnowledgeBase) (bool, error) {
+	compName := fmt.Sprintf("kb-%s", mkb.Name)
+	comp := &v1alpha1.Component{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      compName,
+			Namespace: mkb.Namespace,
+		},
+	}
+	err := r.Get(ctx, client.ObjectKey{Namespace: mkb.Namespace, Name: compName}, comp)
+
+	tmpl, err := templates.Manager.ResolveTemplate(mkb.Spec.Template, "")
+	if err != nil {
+		apimeta.SetStatusCondition(&mkb.Status.Conditions, metav1.Condition{
+			Type:               "KnowledgeBaseRuntimeResolved",
+			Status:             metav1.ConditionFalse,
+			Reason:             "TemplateNotResolved",
+			Message:            fmt.Sprintf("Failed to reconcile template for ManagedKnowledgeBase: %v", err),
+			ObservedGeneration: mkb.Generation,
+		})
+		return false, nil
+	}
+	apimeta.SetStatusCondition(&mkb.Status.Conditions, metav1.Condition{
+		Type:               "KnowledgeBaseRuntimeResolved",
+		Status:             metav1.ConditionTrue,
+		Reason:             "TemplateResolved",
+		Message:            fmt.Sprintf("Successfully resolved template for ManagedKnowledgeBase: %v", tmpl.Name),
+		ObservedGeneration: mkb.Generation,
+	})
+
+	_, err = controllerutil.CreateOrUpdate(ctx, r.Client, comp, func() error {
+		if comp.Labels == nil {
+			comp.Labels = make(map[string]string)
+		}
+		stackName := ""
+		if mkb.Spec.StackRef != nil {
+			stackName = mkb.Spec.StackRef.Name
+		}
+		maps.Copy(comp.Labels, reconciler.ResourceLabels(stackName, "KnowledgeBase", mkb.Name, ""))
+		dependencies := make([]corev1.ObjectReference, 0)
+		if mkb.Spec.EmbeddingModelRef != nil {
+			dependencies = append(dependencies, *mkb.Spec.EmbeddingModelRef)
+		}
+		if mkb.Spec.RerankerModelRef != nil {
+			dependencies = append(dependencies, *mkb.Spec.RerankerModelRef)
+		}
+		comp.Spec.Category = "KnowledgeBase"
+		comp.Spec.Template = &tmpl
+		comp.Spec.Dependencies = dependencies
+		if mkb.Spec.StackRef != nil {
+			comp.Spec.StackRef = &corev1.LocalObjectReference{Name: mkb.Spec.StackRef.Name}
+		}
+		return controllerutil.SetControllerReference(mkb, comp, r.Scheme)
+	})
+
+	available := apimeta.IsStatusConditionTrue(comp.Status.Conditions, "Available")
+	if available {
+		mkb.Status.Endpoint = comp.Status.Endpoint
+		apimeta.SetStatusCondition(&mkb.Status.Conditions, metav1.Condition{
+			Type:               "Ready",
+			Status:             metav1.ConditionTrue,
+			Reason:             "ComponentAvailable",
+			Message:            "Component is available",
+			ObservedGeneration: mkb.Generation,
+		})
+	} else {
+		apimeta.SetStatusCondition(&mkb.Status.Conditions, metav1.Condition{
+			Type:               "Ready",
+			Status:             metav1.ConditionFalse,
+			Reason:             "ComponentUnavailable",
+			Message:            "Component is not available",
+			ObservedGeneration: mkb.Generation,
+		})
+	}
+	return available, err
+}
+
 // SetupWithManager sets up the controller with the Manager.
 func (r *ManagedKnowledgeBaseReconciler) SetupWithManager(mgr ctrl.Manager) error {
+	mapComponentToManagedKnowledgeBases := func(ctx context.Context, obj client.Object) []reconcile.Request {
+		comp, ok := obj.(*v1alpha1.Component)
+		if !ok {
+			return nil
+		}
+		var reqs []reconcile.Request
+		for _, em := range comp.Spec.Dependencies {
+			if em.APIVersion == v1alpha1.GroupVersion.Group && em.Kind == "ManagedKnowledgeBase" {
+				reqs = append(reqs, reconcile.Request{
+					NamespacedName: types.NamespacedName{
+						Name:      em.Name,
+						Namespace: em.Namespace,
+					},
+				})
+			}
+		}
+		return reqs
+	}
+
 	return ctrl.NewControllerManagedBy(mgr).
 		For(&v1alpha1.ManagedKnowledgeBase{}).
+		Owns(&v1alpha1.Component{}).
+		Watches(&v1alpha1.Component{}, handler.EnqueueRequestsFromMapFunc(mapComponentToManagedKnowledgeBases)).
 		Named("managedknowledgebase").
 		Complete(r)
 }