aboutsummaryrefslogtreecommitdiffstats
path: root/vespajlib/src/main/java/ai/vespa/llm/test/MockLanguageModel.java
diff options
context:
space:
mode:
authorJon Bratseth <bratseth@gmail.com>2023-04-26 14:38:13 +0200
committerGitHub <noreply@github.com>2023-04-26 14:38:13 +0200
commit072acd73475aae1a882d908d57d28eec31e69c3e (patch)
tree33ec28c28ff9e2d0badf68df94fff9ec05df5c04 /vespajlib/src/main/java/ai/vespa/llm/test/MockLanguageModel.java
parent5b290735b485e20216fac0733c0d44274eec3937 (diff)
parent901f992bf024fe1b994385ca76b7b59fe2d8a78f (diff)
Merge pull request #26851 from vespa-engine/bratseth/openai-component
Put the openai client in a separate component
Diffstat (limited to 'vespajlib/src/main/java/ai/vespa/llm/test/MockLanguageModel.java')
-rw-r--r--vespajlib/src/main/java/ai/vespa/llm/test/MockLanguageModel.java44
1 files changed, 44 insertions, 0 deletions
diff --git a/vespajlib/src/main/java/ai/vespa/llm/test/MockLanguageModel.java b/vespajlib/src/main/java/ai/vespa/llm/test/MockLanguageModel.java
new file mode 100644
index 00000000000..d47f43c55b2
--- /dev/null
+++ b/vespajlib/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.LanguageModel;
+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.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); }
+
+ }
+
+}