com.yahoo.vespa.model.builder.xml.dom.DomComponentBuilder Maven / Gradle / Ivy

// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.vespa.model.builder.xml.dom;

import com.yahoo.component.ComponentId;
import com.yahoo.config.model.deploy.DeployState;
import com.yahoo.config.model.producer.AnyConfigProducer;
import com.yahoo.config.model.producer.TreeConfigProducer;
import com.yahoo.osgi.provider.model.ComponentModel;
import com.yahoo.text.XML;
import com.yahoo.vespa.model.container.ApplicationContainerCluster;
import com.yahoo.vespa.model.container.component.BertEmbedder;
import com.yahoo.vespa.model.container.component.ColBertEmbedder;
import com.yahoo.vespa.model.container.component.SpladeEmbedder;
import com.yahoo.vespa.model.container.component.Component;
import com.yahoo.vespa.model.container.component.HuggingFaceEmbedder;
import com.yahoo.vespa.model.container.component.HuggingFaceTokenizer;
import com.yahoo.vespa.model.container.xml.BundleInstantiationSpecificationBuilder;
import org.w3c.dom.Element;

/**
 * @author gjoranv
 * @author Tony Vaagenes
 */
public class DomComponentBuilder extends VespaDomBuilder.DomConfigProducerBuilder, AnyConfigProducer> {

    public static final String elementName = "component" ;

    private final ComponentId namespace;

    public DomComponentBuilder() {
        this(null);
    }

    public DomComponentBuilder(ComponentId namespace) {
        this.namespace = namespace;
    }

    @Override
    protected Component, ?> doBuild(DeployState deployState, TreeConfigProducer ancestor, Element spec) {
        var component = buildComponent(spec, deployState, ancestor);
        addChildren(deployState, ancestor, spec, component);
        return component;
    }

    private Component, ?> buildComponent(
            Element spec, DeployState state, TreeConfigProducer ancestor) {
        if (spec.hasAttribute("type")) {
            var type = spec.getAttribute("type");
            return switch (type) {
                case "hugging-face-embedder" -> new HuggingFaceEmbedder((ApplicationContainerCluster)ancestor, spec, state);
                case "hugging-face-tokenizer" -> new HuggingFaceTokenizer(spec, state);
                case "colbert-embedder" -> new ColBertEmbedder((ApplicationContainerCluster)ancestor, spec, state);
                case "bert-embedder" -> new BertEmbedder((ApplicationContainerCluster)ancestor, spec, state);
                case "splade-embedder" -> new SpladeEmbedder((ApplicationContainerCluster)ancestor, spec, state);
                default -> throw new IllegalArgumentException("Unknown component type '%s'".formatted(type));
            };
        } else {
            var bundleSpec = BundleInstantiationSpecificationBuilder.build(spec).nestInNamespace(namespace);
            return new Component<>(new ComponentModel(bundleSpec));
        }
    }

    public static void addChildren(DeployState deployState, TreeConfigProducer ancestor, Element componentNode, Component, ?> component) {
        for (Element childNode : XML.getChildren(componentNode, elementName)) {
            addAndInjectChild(deployState, ancestor, component, childNode);
        }
    }

    private static void addAndInjectChild(DeployState deployState, TreeConfigProducer ancestor, Component, ?> component, Element childNode) {
        Component child = new DomComponentBuilder(component.getComponentId()).build(deployState, ancestor, childNode);
        component.addComponent(child);
        component.inject(child);
    }

}