From eaccd0e3918016013f426eecdb7e0fc3493cece9 Mon Sep 17 00:00:00 2001 From: Jon Bratseth Date: Wed, 19 Apr 2023 11:21:47 +0200 Subject: Avoid ignored directory name --- .../java/ai/vespa/llm/test/MockLanguageModel.java | 44 ++++++++++++++++++++++ .../src/test/java/ai/vespa/llm/CompletionTest.java | 3 +- 2 files changed, 45 insertions(+), 2 deletions(-) create mode 100644 model-integration/src/main/java/ai/vespa/llm/test/MockLanguageModel.java diff --git a/model-integration/src/main/java/ai/vespa/llm/test/MockLanguageModel.java b/model-integration/src/main/java/ai/vespa/llm/test/MockLanguageModel.java new file mode 100644 index 00000000000..54b085a451c --- /dev/null +++ b/model-integration/src/main/java/ai/vespa/llm/test/MockLanguageModel.java @@ -0,0 +1,44 @@ +// Copyright Yahoo. 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.Completion; +import ai.vespa.llm.LanguageModel; +import ai.vespa.llm.Prompt; +import com.yahoo.api.annotations.Beta; + +import java.util.List; +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) { + return completer.apply(prompt); + } + + 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); } + + } + +} diff --git a/model-integration/src/test/java/ai/vespa/llm/CompletionTest.java b/model-integration/src/test/java/ai/vespa/llm/CompletionTest.java index acf7d14c438..30b1c8c2fb1 100644 --- a/model-integration/src/test/java/ai/vespa/llm/CompletionTest.java +++ b/model-integration/src/test/java/ai/vespa/llm/CompletionTest.java @@ -1,7 +1,6 @@ package ai.vespa.llm; -import ai.vespa.llm.testing.MockLanguageModel; -import org.junit.jupiter.api.Assertions; +import ai.vespa.llm.test.MockLanguageModel; import org.junit.jupiter.api.Test; import java.util.List; -- cgit v1.2.3