ai.vespa.llm.test.MockLanguageModel Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of vespajlib Show documentation

Library for use in Java components of Vespa. Shared code which do not fit anywhere else.

There is a newer version: 8.441.21

// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package ai.vespa.llm.test;

import ai.vespa.llm.LanguageModel;
import ai.vespa.llm.InferenceParameters;
import ai.vespa.llm.completion.Completion;
import ai.vespa.llm.completion.Prompt;
import com.yahoo.api.annotations.Beta;

import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;
import java.util.function.Function;

/**
 * @author bratseth
 */
@Beta
public class MockLanguageModel implements LanguageModel {

    private final Function> completer;

    public MockLanguageModel(Builder builder) {
        completer = builder.completer;
    }

    @Override
    public List complete(Prompt prompt, InferenceParameters options) {
        return completer.apply(prompt);
    }

    @Override
    public CompletableFuture completeAsync(Prompt prompt,
                                                                    InferenceParameters options,
                                                                    Consumer action) {
        throw new RuntimeException("Not implemented");
    }

    public static class Builder {

        private Function> completer = prompt -> List.of(Completion.from(""));

        public Builder completer(Function> completer) {
            this.completer = completer;
            return this;
        }

        public Builder() {}

        public MockLanguageModel build() { return new MockLanguageModel(this); }

    }

}