diff options
author | Jon Bratseth <bratseth@vespa.ai> | 2023-04-19 11:21:47 +0200 |
---|---|---|
committer | Jon Bratseth <bratseth@vespa.ai> | 2023-04-19 11:21:47 +0200 |
commit | eaccd0e3918016013f426eecdb7e0fc3493cece9 (patch) | |
tree | 0e209fb66c44bb29a9125e1a3f16cf3d16eceea6 | |
parent | 2d9fd2e6c78e9ae8580cd4a21d20d0febb8f9c93 (diff) |
Avoid ignored directory name
-rw-r--r-- | model-integration/src/main/java/ai/vespa/llm/test/MockLanguageModel.java | 44 | ||||
-rw-r--r-- | model-integration/src/test/java/ai/vespa/llm/CompletionTest.java | 3 |
2 files changed, 45 insertions, 2 deletions
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<Prompt, List<Completion>> completer; + + public MockLanguageModel(Builder builder) { + completer = builder.completer; + } + + @Override + public List<Completion> complete(Prompt prompt) { + return completer.apply(prompt); + } + + public static class Builder { + + private Function<Prompt, List<Completion>> completer = prompt -> List.of(Completion.from("")); + + public Builder completer(Function<Prompt, List<Completion>> 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; |