Browse Source

feat(agent): implement Agent controller

Thomas Zhang 4 months ago
parent
commit
df11cb0c66
2 changed files with 396 additions and 11 deletions
  1. 138 11
      internal/controller/agent_controller.go
  2. 258 0
      internal/reconciler/agent.go

+ 138 - 11
internal/controller/agent_controller.go

@@ -18,13 +18,22 @@ package controller
 
 
 import (
 import (
 	"context"
 	"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/runtime"
 	ctrl "sigs.k8s.io/controller-runtime"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
 	logf "sigs.k8s.io/controller-runtime/pkg/log"
 	logf "sigs.k8s.io/controller-runtime/pkg/log"
 
 
 	"github.com/LocoStack/loco-operator/api/v1alpha1"
 	"github.com/LocoStack/loco-operator/api/v1alpha1"
+	"github.com/LocoStack/loco-operator/internal/reconciler"
+	"github.com/LocoStack/loco-operator/pkg/templates"
 )
 )
 
 
 // AgentReconciler reconciles a Agent object
 // AgentReconciler reconciles a Agent object
@@ -36,28 +45,146 @@ type AgentReconciler struct {
 // +kubebuilder:rbac:groups=locostack.com,resources=agents,verbs=get;list;watch;create;update;patch;delete
 // +kubebuilder:rbac:groups=locostack.com,resources=agents,verbs=get;list;watch;create;update;patch;delete
 // +kubebuilder:rbac:groups=locostack.com,resources=agents/status,verbs=get;update;patch
 // +kubebuilder:rbac:groups=locostack.com,resources=agents/status,verbs=get;update;patch
 // +kubebuilder:rbac:groups=locostack.com,resources=agents/finalizers,verbs=update
 // +kubebuilder:rbac:groups=locostack.com,resources=agents/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 Agent 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 *AgentReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
 func (r *AgentReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
-	_ = logf.FromContext(ctx)
+	log := logf.FromContext(ctx)
 
 
-	// TODO(user): your logic here
+	agent := &v1alpha1.Agent{}
+	if err := r.Get(ctx, req.NamespacedName, agent); err != nil {
+		if apierrors.IsNotFound(err) {
+			return ctrl.Result{}, nil
+		}
+		log.Error(err, "Failed to get Agent", "namespace", req.NamespacedName.Namespace, "name", req.NamespacedName.Name)
+		return ctrl.Result{}, err
+	}
 
 
+	patch := client.MergeFrom(agent.DeepCopy())
+
+	componentAvailable, err := r.reconcileComponent(ctx, agent)
+	if err != nil {
+		return ctrl.Result{}, err
+	}
+	if !componentAvailable {
+		apimeta.SetStatusCondition(&agent.Status.Conditions, metav1.Condition{
+			Type:               "Available",
+			Status:             metav1.ConditionFalse,
+			Reason:             "AgentUnavailable",
+			Message:            "Agent is not available",
+			ObservedGeneration: agent.Generation,
+		})
+	} else {
+		apimeta.SetStatusCondition(&agent.Status.Conditions, metav1.Condition{
+			Type:               "Available",
+			Status:             metav1.ConditionTrue,
+			Reason:             "AgentAvailable",
+			Message:            "Agent is available",
+			ObservedGeneration: agent.Generation,
+		})
+	}
+
+	agent.Status.ObservedGeneration = agent.Generation
+	if err := r.Status().Patch(ctx, agent, patch); err != nil {
+		return ctrl.Result{}, client.IgnoreNotFound(err)
+	}
+
+	log.Info("Reconciled Agent", "namespace", req.NamespacedName.Namespace, "name", req.NamespacedName.Name)
 	return ctrl.Result{}, nil
 	return ctrl.Result{}, nil
 }
 }
 
 
+func (r *AgentReconciler) reconcileComponent(ctx context.Context, agent *v1alpha1.Agent) (bool, error) {
+	compName := fmt.Sprintf("agent-%s", agent.Name)
+	comp := &v1alpha1.Component{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      compName,
+			Namespace: agent.Namespace,
+		},
+	}
+	err := r.Get(ctx, client.ObjectKey{Namespace: agent.Namespace, Name: compName}, comp)
+
+	tmpl, err := templates.Manager.ResolveTemplate(agent.Spec.Template, "")
+	if err != nil {
+		apimeta.SetStatusCondition(&agent.Status.Conditions, metav1.Condition{
+			Type:               "AgentRuntimeResolved",
+			Status:             metav1.ConditionFalse,
+			Reason:             "TemplateNotResolved",
+			Message:            fmt.Sprintf("Failed to reconcile template for Agent: %v", err),
+			ObservedGeneration: agent.Generation,
+		})
+		return false, nil
+	}
+	apimeta.SetStatusCondition(&agent.Status.Conditions, metav1.Condition{
+		Type:               "AgentRuntimeResolved",
+		Status:             metav1.ConditionTrue,
+		Reason:             "TemplateResolved",
+		Message:            fmt.Sprintf("Successfully resolved template for Agent: %v", tmpl.Name),
+		ObservedGeneration: agent.Generation,
+	})
+
+	_, err = controllerutil.CreateOrUpdate(ctx, r.Client, comp, func() error {
+		if comp.Labels == nil {
+			comp.Labels = make(map[string]string)
+		}
+		stackName := ""
+		if agent.Spec.StackRef != nil {
+			stackName = agent.Spec.StackRef.Name
+		}
+		maps.Copy(comp.Labels, reconciler.ResourceLabels(stackName, "Agent", agent.Name, ""))
+		dependencies := make([]corev1.ObjectReference, 0)
+		dependencies = append(dependencies, agent.Spec.Model)
+		for _, skill := range agent.Spec.Skills {
+			dependencies = append(dependencies, skill)
+		}
+		for _, tool := range agent.Spec.Tools {
+			dependencies = append(dependencies, tool)
+		}
+		for _, kb := range agent.Spec.KnowledgeBases {
+			dependencies = append(dependencies, kb)
+		}
+		for _, subagent := range agent.Spec.Subagents {
+			dependencies = append(dependencies, subagent)
+		}
+		comp.Spec.Category = "Agent"
+		comp.Spec.Template = &tmpl
+		comp.Spec.Variables = map[string]string{
+			"systemPrompt": agent.Spec.SystemPrompt,
+		}
+		comp.Spec.Dependencies = dependencies
+		if agent.Spec.StackRef != nil {
+			comp.Spec.StackRef = &corev1.LocalObjectReference{Name: agent.Spec.StackRef.Name}
+		}
+		return controllerutil.SetControllerReference(agent, comp, r.Scheme)
+	})
+
+	available := apimeta.IsStatusConditionTrue(comp.Status.Conditions, "Available")
+	if available {
+		agent.Status.Endpoint = comp.Status.Endpoint
+		apimeta.SetStatusCondition(&agent.Status.Conditions, metav1.Condition{
+			Type:               "Ready",
+			Status:             metav1.ConditionTrue,
+			Reason:             "ComponentAvailable",
+			Message:            "Component is available",
+			ObservedGeneration: agent.Generation,
+		})
+	} else {
+		apimeta.SetStatusCondition(&agent.Status.Conditions, metav1.Condition{
+			Type:               "Ready",
+			Status:             metav1.ConditionFalse,
+			Reason:             "ComponentUnavailable",
+			Message:            "Component is not available",
+			ObservedGeneration: agent.Generation,
+		})
+	}
+	return available, err
+}
+
 // SetupWithManager sets up the controller with the Manager.
 // SetupWithManager sets up the controller with the Manager.
 func (r *AgentReconciler) SetupWithManager(mgr ctrl.Manager) error {
 func (r *AgentReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	return ctrl.NewControllerManagedBy(mgr).
 	return ctrl.NewControllerManagedBy(mgr).
 		For(&v1alpha1.Agent{}).
 		For(&v1alpha1.Agent{}).
+		Owns(&v1alpha1.Component{}).
+		Watches(&v1alpha1.Component{}, ComponentEventHandler("Agent")).
 		Named("agent").
 		Named("agent").
 		Complete(r)
 		Complete(r)
 }
 }

+ 258 - 0
internal/reconciler/agent.go

@@ -0,0 +1,258 @@
+package reconciler
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"maps"
+	"sort"
+	"strings"
+
+	"github.com/LocoStack/loco-operator/api/v1alpha1"
+	"github.com/LocoStack/loco-operator/pkg/templates/litellm"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
+)
+
+type AgentReconciler struct {
+	client    client.Client
+	scheme    *runtime.Scheme
+	stack     *v1alpha1.Stack
+	component *v1alpha1.Component
+}
+
+func NewAgentReconciler(client client.Client, scheme *runtime.Scheme, stack *v1alpha1.Stack, component *v1alpha1.Component) *AgentReconciler {
+	return &AgentReconciler{
+		client:    client,
+		scheme:    scheme,
+		stack:     stack,
+		component: component,
+	}
+}
+
+type agentComponentEndpoint struct {
+	Name string
+	URL  string
+}
+
+type AgentComponents struct {
+	Gateway        agentComponentEndpoint
+	GatewayKey     string
+	Model          agentComponentEndpoint
+	Skills         []string
+	Tools          []agentComponentEndpoint
+	KnowledgeBases []agentComponentEndpoint
+	Agents         []agentComponentEndpoint
+}
+
+func (r *AgentReconciler) getDependencies(ctx context.Context) ([]client.Object, *AgentComponents, error) {
+	comps := &AgentComponents{
+		Tools:          []agentComponentEndpoint{},
+		KnowledgeBases: []agentComponentEndpoint{},
+	}
+	endpointResolver := NewEndpointResolver(r.client, r.scheme, r.stack.Namespace)
+	deps := make([]client.Object, 0)
+	for _, dep := range r.component.Spec.Dependencies {
+		obj, endpoint, err := endpointResolver.ResolveEndpoint(ctx, r.stack, dep)
+		if err != nil {
+			return nil, nil, fmt.Errorf("Failed to resolve endpoint for %s %s: %w", dep.Kind, dep.Name, err)
+		}
+		switch dep.Kind {
+		case "ExternalModel":
+			fallthrough
+		case "ManagedModel":
+			var category string
+			if extModel, ok := obj.(*v1alpha1.ExternalModel); ok {
+				category = extModel.Spec.Category
+				comps.Model.Name = extModel.Spec.ModelName
+			} else if managedModel, ok := obj.(*v1alpha1.ManagedModel); ok {
+				category = managedModel.Spec.Category
+				comps.Model.Name = managedModel.Spec.ModelName
+			} else {
+				return nil, nil, fmt.Errorf("Unsupported dependency kind for agent: %s", dep.Kind)
+			}
+			if category != "language" {
+				return nil, nil, fmt.Errorf("Unsupported model category for agent: %s", category)
+			}
+			comps.Model.URL = endpoint
+		case "Skill":
+			comps.Skills = append(comps.Skills, dep.Name)
+		case "ExternalTool":
+			fallthrough
+		case "ManagedTool":
+			var toolName string
+			if extTool, ok := obj.(*v1alpha1.ExternalTool); ok {
+				toolName = extTool.Spec.ToolName
+			} else if managedTool, ok := obj.(*v1alpha1.ManagedTool); ok {
+				toolName = managedTool.Spec.ToolName
+			}
+			comps.Tools = append(comps.Tools, agentComponentEndpoint{Name: toolName, URL: endpoint})
+		case "ExternalKnowledgeBase":
+			fallthrough
+		case "ManagedKnowledgeBase":
+			comps.KnowledgeBases = append(comps.KnowledgeBases, agentComponentEndpoint{Name: dep.Name, URL: endpoint})
+		case "Agent":
+			comps.Agents = append(comps.Agents, agentComponentEndpoint{Name: dep.Name, URL: endpoint})
+		default:
+			return nil, nil, fmt.Errorf("Unsupported dependency kind for agent: %s", dep.Kind)
+		}
+		deps = append(deps, obj)
+	}
+	sort.Strings(comps.Skills)
+	sort.Slice(comps.Tools, func(i, j int) bool {
+		return comps.Tools[i].Name < comps.Tools[j].Name
+	})
+	sort.Slice(comps.KnowledgeBases, func(i, j int) bool {
+		return comps.KnowledgeBases[i].Name < comps.KnowledgeBases[j].Name
+	})
+	sort.Slice(comps.Agents, func(i, j int) bool {
+		return comps.Agents[i].Name < comps.Agents[j].Name
+	})
+
+	gatewayComp, err := endpointResolver.resolveGatewayEndpoint(ctx, r.stack)
+	if err != nil {
+		return nil, nil, fmt.Errorf("Failed to resolve gateway for %s: %w", r.component.Name, err)
+	}
+	comps.Gateway = agentComponentEndpoint{Name: gatewayComp.Name, URL: gatewayComp.Status.Endpoint}
+	masterKey, err := r.gatewayMasterKey(ctx, gatewayComp)
+	if err != nil {
+		return nil, nil, fmt.Errorf("resolving gateway master key: %w", err)
+	}
+	comps.GatewayKey = masterKey
+	return deps, comps, nil
+}
+
+func (r *AgentReconciler) gatewayMasterKey(ctx context.Context, gatewayComp *v1alpha1.Component) (string, error) {
+	if gatewayComp == nil {
+		return "", nil
+	}
+	secret := &corev1.Secret{}
+	secretName := litellm.MasterKeySecretName(gatewayComp.Name)
+	if err := r.client.Get(ctx, client.ObjectKey{Name: secretName, Namespace: r.component.GetNamespace()}, secret); err != nil {
+		return "", fmt.Errorf("Failed to get gateway master key secret %s: %w", secretName, err)
+	}
+	key, ok := secret.Data[litellm.LITELLM_AUTH_SECRET_KEY]
+	if !ok {
+		return "", fmt.Errorf("Secret %s does not contain key %q", secretName, litellm.LITELLM_AUTH_SECRET_KEY)
+	}
+	return string(key), nil
+}
+
+// AddSkillsToTmpl adds skill mounts to the template,
+// mounting each dependency Skill's ConfigMap under <skillsRoot>/skills/<skillName>,
+// and returns a hash of the mounted files. Agent runtimes read their skills once at
+// startup, and a ConfigMap edited in place never restarts the Pod, so callers fold that
+// hash into the Pod template to roll the agent out whenever a Skill's content changes.
+// A Skill's ConfigMap stores one entry per file (SKILL.md plus its scripts/references/
+// assets), keyed by v1alpha1.SkillConfigMapKey(path) rather than the raw relative path,
+// since ConfigMap keys cannot contain "/". The ConfigMap's v1alpha1.SkillFilePathsAnnotation
+// records the key -> original relative path mapping, so it can be read directly rather than
+// re-fetched and recomputed from the Skill spec; a ConfigMapVolumeSource with per-key Items
+// then lays the files back out as a real directory tree using that mapping.
+func (r *AgentReconciler) AddSkillsToTmpl(ctx context.Context, comps *AgentComponents, tmpl *v1alpha1.Template, skillsRoot string) (string, error) {
+	volumes := make([]corev1.Volume, 0, len(comps.Skills))
+	volumeMounts := make([]corev1.VolumeMount, 0, len(comps.Skills))
+	skillsHash := sha256.New()
+	for _, skillName := range comps.Skills {
+		cm := &corev1.ConfigMap{}
+		if err := r.client.Get(ctx, client.ObjectKey{Namespace: r.stack.Namespace, Name: skillName}, cm); err != nil {
+			return "", fmt.Errorf("Failed to get ConfigMap for Skill %s: %w", skillName, err)
+		}
+		filePathsJSON, ok := cm.Annotations[v1alpha1.SkillFilePathsAnnotation]
+		if !ok {
+			return "", fmt.Errorf("ConfigMap for Skill %s is missing the %q annotation", skillName, v1alpha1.SkillFilePathsAnnotation)
+		}
+		var filePaths map[string]string
+		if err := json.Unmarshal([]byte(filePathsJSON), &filePaths); err != nil {
+			return "", fmt.Errorf("unmarshaling file paths for Skill %s: %w", skillName, err)
+		}
+
+		keys := sortedKeys(filePaths)
+
+		items := make([]corev1.KeyToPath, 0, len(keys))
+		for _, key := range keys {
+			items = append(items, corev1.KeyToPath{
+				Key:  key,
+				Path: filePaths[key],
+			})
+			_, _ = fmt.Fprintf(skillsHash, "%s/%s=%s\n", skillName, filePaths[key], cm.Data[key])
+		}
+		volumes = append(volumes, corev1.Volume{
+			Name: skillName,
+			VolumeSource: corev1.VolumeSource{
+				ConfigMap: &corev1.ConfigMapVolumeSource{
+					LocalObjectReference: corev1.LocalObjectReference{Name: skillName},
+					Items:                items,
+				},
+			},
+		})
+		volumeMounts = append(volumeMounts, corev1.VolumeMount{
+			Name:      skillName,
+			MountPath: fmt.Sprintf("%s/skills/%s", skillsRoot, skillName),
+			ReadOnly:  true,
+		})
+	}
+	tmpl.Spec.Volumes = append(tmpl.Spec.Volumes, volumes...)
+	tmpl.Spec.Runtime.VolumeMounts = append(tmpl.Spec.Runtime.VolumeMounts, volumeMounts...)
+	return hex.EncodeToString(skillsHash.Sum(nil)), nil
+}
+
+type mcpConfig struct {
+	MCPServers map[string]mcpServer `json:"mcpServers"`
+}
+
+type mcpServer struct {
+	URL     string            `json:"url"`
+	Headers map[string]string `json:"headers,omitempty"`
+}
+
+// reconcileMcpConfig writes the agent's mcp.json as a Secret (not a ConfigMap): when a
+// gateway is present, the generated config embeds the resolved litellm master key value
+// in the auth header, so the credential must not land in a ConfigMap.
+func (r *AgentReconciler) reconcileMcpConfig(ctx context.Context, comps *AgentComponents) (string, error) {
+	cfg := mcpConfig{MCPServers: make(map[string]mcpServer)}
+	if comps.Gateway.URL != "" && len(comps.Tools) > 0 {
+		toolNames := make([]string, 0, len(comps.Tools))
+		for _, tool := range comps.Tools {
+			toolNames = append(toolNames, tool.Name)
+		}
+		sort.Strings(toolNames)
+		cfg.MCPServers["litellm"] = mcpServer{
+			URL: fmt.Sprintf("%s/%s/mcp", comps.Gateway.URL, strings.Join(toolNames, ",")),
+			Headers: map[string]string{
+				litellm.MCPAuthHeaderName: "Bearer " + comps.GatewayKey,
+			},
+		}
+	} else {
+		for _, tool := range comps.Tools {
+			cfg.MCPServers[tool.Name] = mcpServer{URL: tool.URL}
+		}
+	}
+	jsonBytes, err := json.Marshal(cfg)
+	if err != nil {
+		return "", fmt.Errorf("marshaling agent mcp config: %w", err)
+	}
+	h := sha256.Sum256(jsonBytes)
+	configHash := hex.EncodeToString(h[:])
+	secret := &corev1.Secret{ObjectMeta: metav1.ObjectMeta{Name: fmt.Sprintf("%s-mcp-config", r.component.GetName()), Namespace: r.component.GetNamespace()}}
+	_, err = controllerutil.CreateOrUpdate(ctx, r.client, secret, func() error {
+		if secret.Labels == nil {
+			secret.Labels = map[string]string{}
+		}
+		maps.Copy(secret.Labels, ResourceLabels(r.stack.Name, r.component.Spec.Category, r.component.GetName(), ""))
+		if secret.Annotations == nil {
+			secret.Annotations = make(map[string]string)
+		}
+		secret.Annotations["hash"] = configHash
+		secret.Data = map[string][]byte{
+			"mcp.json": jsonBytes,
+		}
+		return controllerutil.SetControllerReference(r.component, secret, r.scheme)
+	})
+	return configHash, err
+}