feat: add endpoint resolver

internal/reconciler/common.go

@@ -23,6 +23,7 @@ import (
 
 	"github.com/LocoStack/loco-operator/api/v1alpha1"
 	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
@@ -173,3 +174,148 @@ func buildPodSpec(variables map[string]string, tmpl *v1alpha1.Template) (corev1.
 		Volumes:                   buildVolumes(variables, tmpl),
 	}, nil
 }
+
+type EndpointResolver struct {
+	client    client.Client
+	scheme    *runtime.Scheme
+	namespace string
+}
+
+func NewEndpointResolver(client client.Client, scheme *runtime.Scheme, namespace string) *EndpointResolver {
+	return &EndpointResolver{
+		client:    client,
+		scheme:    scheme,
+		namespace: namespace,
+	}
+}
+
+func (r *EndpointResolver) ResolveEndpoint(ctx context.Context, stack *v1alpha1.Stack, ref corev1.ObjectReference) (client.Object, string, error) {
+	switch ref.Kind {
+	case "ExternalModel":
+		return r.resolveExternalModelEndpoint(ctx, stack, ref.Name)
+	case "ManagedModel":
+		return r.resolveManagedModelEndpoint(ctx, stack, ref.Name)
+	case "ExternalTool":
+		return r.resolveExternalToolEndpoint(ctx, stack, ref.Name)
+	case "ManagedTool":
+		return r.resolveManagedToolEndpoint(ctx, stack, ref.Name)
+	case "ExternalKnowledgeBase":
+		return r.resolveExternalKnowledgeBaseEndpoint(ctx, ref.Name)
+	case "ManagedKnowledgeBase":
+		return r.resolveManagedKnowledgeBaseEndpoint(ctx, ref.Name)
+	case "Component":
+		return r.resolveComponentEndpoint(ctx, ref.Name)
+	default:
+		return nil, "", fmt.Errorf("unsupported reference kind: %s", ref.Kind)
+	}
+}
+
+func (r *EndpointResolver) resolveGatewayEndpoint(ctx context.Context, stack *v1alpha1.Stack) (*v1alpha1.Component, error) {
+	gatewayEnabled := true
+	if stack.Spec.Gateway != nil {
+		gatewayEnabled = stack.Spec.Gateway.Enabled
+	}
+	if !gatewayEnabled {
+		return nil, nil
+	}
+	gwList := &v1alpha1.ComponentList{}
+	err := r.client.List(ctx, gwList, client.InNamespace(stack.Namespace), client.MatchingLabels{"locostack.com/component": "Gateway", "stack.locostack.com/default": "true"})
+	if err != nil {
+		return nil, fmt.Errorf("Failed to list gateway components: %v", err)
+	}
+	if len(gwList.Items) == 0 {
+		return nil, fmt.Errorf("No gateway component found for stack %s", stack.Name)
+	} else if len(gwList.Items) > 1 {
+		return nil, fmt.Errorf("Multiple gateway components found for stack %s", stack.Name)
+	}
+	gatewayComp := &gwList.Items[0]
+	if gatewayComp.Status.Endpoint == "" {
+		return nil, fmt.Errorf("Gateway component for stack %s not ready", stack.Name)
+	}
+	return gatewayComp, nil
+}
+
+func (r *EndpointResolver) resolveExternalModelEndpoint(ctx context.Context, stack *v1alpha1.Stack, name string) (client.Object, string, error) {
+	var model v1alpha1.ExternalModel
+	if err := r.client.Get(ctx, client.ObjectKey{Namespace: r.namespace, Name: name}, &model); err != nil {
+		return nil, "", fmt.Errorf("Failed to get ExternalModel %s: %w", name, err)
+	}
+	gatewayComp, err := r.resolveGatewayEndpoint(ctx, stack)
+	if err != nil {
+		return nil, "", fmt.Errorf("Failed to resolve gateway endpoint for ExternalModel %s: %w", name, err)
+	}
+	if gatewayComp != nil {
+		return &model, gatewayComp.Status.Endpoint, nil
+	}
+	return &model, model.Spec.APIBase, nil
+}
+
+func (r *EndpointResolver) resolveManagedModelEndpoint(ctx context.Context, stack *v1alpha1.Stack, name string) (client.Object, string, error) {
+	var model v1alpha1.ManagedModel
+	if err := r.client.Get(ctx, client.ObjectKey{Namespace: r.namespace, Name: name}, &model); err != nil {
+		return nil, "", fmt.Errorf("Failed to get ManagedModel %s: %w", name, err)
+	}
+	gatewayComp, err := r.resolveGatewayEndpoint(ctx, stack)
+	if err != nil {
+		return nil, "", fmt.Errorf("Failed to resolve gateway endpoint for ManagedModel %s: %w", name, err)
+	}
+	if gatewayComp != nil {
+		return &model, gatewayComp.Status.Endpoint + "/v1", nil
+	}
+	return &model, model.Status.Endpoint, nil
+}
+
+func (r *EndpointResolver) resolveExternalToolEndpoint(ctx context.Context, stack *v1alpha1.Stack, name string) (client.Object, string, error) {
+	var tool v1alpha1.ExternalTool
+	if err := r.client.Get(ctx, client.ObjectKey{Namespace: r.namespace, Name: name}, &tool); err != nil {
+		return nil, "", fmt.Errorf("Failed to get ExternalTool %s: %w", name, err)
+	}
+	gatewayComp, err := r.resolveGatewayEndpoint(ctx, stack)
+	if err != nil {
+		return nil, "", fmt.Errorf("Failed to resolve gateway endpoint for ExternalTool %s: %w", name, err)
+	}
+	if gatewayComp != nil {
+		return &tool, gatewayComp.Status.Endpoint + "/mcp", nil
+	}
+	return &tool, tool.Spec.Endpoint, nil
+}
+
+func (r *EndpointResolver) resolveManagedToolEndpoint(ctx context.Context, stack *v1alpha1.Stack, name string) (client.Object, string, error) {
+	var tool v1alpha1.ManagedTool
+	if err := r.client.Get(ctx, client.ObjectKey{Namespace: r.namespace, Name: name}, &tool); err != nil {
+		return nil, "", fmt.Errorf("Failed to get ManagedTool %s: %w", name, err)
+	}
+	gatewayComp, err := r.resolveGatewayEndpoint(ctx, stack)
+	if err != nil {
+		return nil, "", fmt.Errorf("Failed to resolve gateway endpoint for ManagedTool %s: %w", name, err)
+	}
+	if gatewayComp != nil {
+		return &tool, gatewayComp.Status.Endpoint, nil
+	}
+	return &tool, tool.Status.Endpoint, nil
+}
+
+func (r *EndpointResolver) resolveExternalKnowledgeBaseEndpoint(ctx context.Context, name string) (client.Object, string, error) {
+	var kb v1alpha1.ExternalKnowledgeBase
+	if err := r.client.Get(ctx, client.ObjectKey{Namespace: r.namespace, Name: name}, &kb); err != nil {
+		return nil, "", fmt.Errorf("Failed to get ExternalKnowledgeBase %s: %w", name, err)
+	}
+	// TODO: use gateway endpoint
+	return &kb, "", nil
+}
+
+func (r *EndpointResolver) resolveManagedKnowledgeBaseEndpoint(ctx context.Context, name string) (client.Object, string, error) {
+	var kb v1alpha1.ManagedKnowledgeBase
+	if err := r.client.Get(ctx, client.ObjectKey{Namespace: r.namespace, Name: name}, &kb); err != nil {
+		return nil, "", fmt.Errorf("Failed to get ManagedKnowledgeBase %s: %w", name, err)
+	}
+	return &kb, kb.Status.Endpoint, nil
+}
+
+func (r *EndpointResolver) resolveComponentEndpoint(ctx context.Context, name string) (client.Object, string, error) {
+	comp := &v1alpha1.Component{}
+	if err := r.client.Get(ctx, client.ObjectKey{Namespace: r.namespace, Name: name}, comp); err != nil {
+		return nil, "", fmt.Errorf("Failed to get Component %s: %w", name, err)
+	}
+	return comp, comp.Status.Endpoint, nil
+}