Put the openai client in a separate component

10 files changed, 352 insertions, 0 deletions

diff --git a/vespajlib/abi-spec.json b/vespajlib/abi-spec.json
index 48d519b3a62..fe07460c20b 100644
--- a/vespajlib/abi-spec.json
+++ b/vespajlib/abi-spec.json
@@ -3995,5 +3995,113 @@
       "public void leaving(com.yahoo.yolean.trace.TraceNode)"
     ],
     "fields" : [ ]
+  },
+  "ai.vespa.llm.LanguageModel" : {
+    "superClass" : "java.lang.Object",
+    "interfaces" : [ ],
+    "attributes" : [
+      "public",
+      "interface",
+      "abstract"
+    ],
+    "methods" : [
+      "public abstract java.util.List complete(ai.vespa.llm.completion.Prompt)"
+    ],
+    "fields" : [ ]
+  },
+  "ai.vespa.llm.completion.Completion$FinishReason" : {
+    "superClass" : "java.lang.Enum",
+    "interfaces" : [ ],
+    "attributes" : [
+      "public",
+      "final",
+      "enum"
+    ],
+    "methods" : [
+      "public static ai.vespa.llm.completion.Completion$FinishReason[] values()",
+      "public static ai.vespa.llm.completion.Completion$FinishReason valueOf(java.lang.String)"
+    ],
+    "fields" : [
+      "public static final enum ai.vespa.llm.completion.Completion$FinishReason length",
+      "public static final enum ai.vespa.llm.completion.Completion$FinishReason stop"
+    ]
+  },
+  "ai.vespa.llm.completion.Completion" : {
+    "superClass" : "java.lang.Record",
+    "interfaces" : [ ],
+    "attributes" : [
+      "public",
+      "final",
+      "record"
+    ],
+    "methods" : [
+      "public void <init>(java.lang.String, ai.vespa.llm.completion.Completion$FinishReason)",
+      "public java.lang.String text()",
+      "public ai.vespa.llm.completion.Completion$FinishReason finishReason()",
+      "public static ai.vespa.llm.completion.Completion from(java.lang.String)",
+      "public final java.lang.String toString()",
+      "public final int hashCode()",
+      "public final boolean equals(java.lang.Object)"
+    ],
+    "fields" : [ ]
+  },
+  "ai.vespa.llm.completion.Prompt" : {
+    "superClass" : "java.lang.Object",
+    "interfaces" : [ ],
+    "attributes" : [
+      "public",
+      "abstract"
+    ],
+    "methods" : [
+      "public void <init>()",
+      "public abstract java.lang.String asString()",
+      "public ai.vespa.llm.completion.Prompt append(ai.vespa.llm.completion.Completion)",
+      "public abstract ai.vespa.llm.completion.Prompt append(java.lang.String)"
+    ],
+    "fields" : [ ]
+  },
+  "ai.vespa.llm.completion.StringPrompt" : {
+    "superClass" : "ai.vespa.llm.completion.Prompt",
+    "interfaces" : [ ],
+    "attributes" : [
+      "public"
+    ],
+    "methods" : [
+      "public java.lang.String asString()",
+      "public ai.vespa.llm.completion.StringPrompt append(java.lang.String)",
+      "public ai.vespa.llm.completion.StringPrompt append(ai.vespa.llm.completion.Completion)",
+      "public java.lang.String toString()",
+      "public static ai.vespa.llm.completion.StringPrompt from(java.lang.String)",
+      "public bridge synthetic ai.vespa.llm.completion.Prompt append(java.lang.String)",
+      "public bridge synthetic ai.vespa.llm.completion.Prompt append(ai.vespa.llm.completion.Completion)"
+    ],
+    "fields" : [ ]
+  },
+  "ai.vespa.llm.test.MockLanguageModel$Builder" : {
+    "superClass" : "java.lang.Object",
+    "interfaces" : [ ],
+    "attributes" : [
+      "public"
+    ],
+    "methods" : [
+      "public ai.vespa.llm.test.MockLanguageModel$Builder completer(java.util.function.Function)",
+      "public void <init>()",
+      "public ai.vespa.llm.test.MockLanguageModel build()"
+    ],
+    "fields" : [ ]
+  },
+  "ai.vespa.llm.test.MockLanguageModel" : {
+    "superClass" : "java.lang.Object",
+    "interfaces" : [
+      "ai.vespa.llm.LanguageModel"
+    ],
+    "attributes" : [
+      "public"
+    ],
+    "methods" : [
+      "public void <init>(ai.vespa.llm.test.MockLanguageModel$Builder)",
+      "public java.util.List complete(ai.vespa.llm.completion.Prompt)"
+    ],
+    "fields" : [ ]
   }
 }
\ No newline at end of file
diff --git a/vespajlib/src/main/java/ai/vespa/llm/LanguageModel.java b/vespajlib/src/main/java/ai/vespa/llm/LanguageModel.java
new file mode 100644
index 00000000000..829b74f7bf4
--- /dev/null
+++ b/vespajlib/src/main/java/ai/vespa/llm/LanguageModel.java
@@ -0,0 +1,20 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package ai.vespa.llm;
+
+import ai.vespa.llm.completion.Completion;
+import ai.vespa.llm.completion.Prompt;
+import com.yahoo.api.annotations.Beta;
+
+import java.util.List;
+
+/**
+ * Interface to language models.
+ *
+ * @author bratseth
+ */
+@Beta
+public interface LanguageModel {
+
+    List<Completion> complete(Prompt prompt);
+
+}
diff --git a/vespajlib/src/main/java/ai/vespa/llm/completion/Completion.java b/vespajlib/src/main/java/ai/vespa/llm/completion/Completion.java
new file mode 100644
index 00000000000..30645b5151f
--- /dev/null
+++ b/vespajlib/src/main/java/ai/vespa/llm/completion/Completion.java
@@ -0,0 +1,41 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package ai.vespa.llm.completion;
+
+import com.yahoo.api.annotations.Beta;
+
+import java.util.Objects;
+
+/**
+ * A completion from a language model.
+ *
+ * @author bratseth
+ */
+@Beta
+public record Completion(String text, FinishReason finishReason) {
+
+    public enum FinishReason {
+
+        /** The maximum length of a completion was reached. */
+        length,
+
+        /** The completion is the predicted ending of the prompt. */
+        stop
+
+    }
+
+    public Completion(String text, FinishReason finishReason) {
+        this.text = Objects.requireNonNull(text);
+        this.finishReason = Objects.requireNonNull(finishReason);
+    }
+
+    /** Returns the generated text completion. */
+    public String text() { return text; }
+
+    /** Returns the reason this completion ended. */
+    public FinishReason finishReason() { return finishReason; }
+
+    public static Completion from(String text) {
+        return new Completion(text, FinishReason.stop);
+    }
+
+}
diff --git a/vespajlib/src/main/java/ai/vespa/llm/completion/Prompt.java b/vespajlib/src/main/java/ai/vespa/llm/completion/Prompt.java
new file mode 100644
index 00000000000..d5d0247d6b0
--- /dev/null
+++ b/vespajlib/src/main/java/ai/vespa/llm/completion/Prompt.java
@@ -0,0 +1,23 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package ai.vespa.llm.completion;
+
+import com.yahoo.api.annotations.Beta;
+
+/**
+ * A prompt that can be given to a large language model to generate a completion.
+ *
+ * @author bratseth
+ */
+@Beta
+public abstract class Prompt {
+
+    public abstract String asString();
+
+    /** Returns a new prompt with the text of the given completion appended. */
+    public Prompt append(Completion completion) {
+        return append(completion.text());
+    }
+
+    public abstract Prompt append(String text);
+
+}
diff --git a/vespajlib/src/main/java/ai/vespa/llm/completion/StringPrompt.java b/vespajlib/src/main/java/ai/vespa/llm/completion/StringPrompt.java
new file mode 100644
index 00000000000..e8392ca992e
--- /dev/null
+++ b/vespajlib/src/main/java/ai/vespa/llm/completion/StringPrompt.java
@@ -0,0 +1,43 @@
+package ai.vespa.llm.completion;
+
+import com.yahoo.api.annotations.Beta;
+
+import java.util.Objects;
+
+/**
+ * A prompt which just consists of a string.
+ *
+ * @author bratseth
+ */
+@Beta
+public class StringPrompt extends Prompt {
+
+    private final String string;
+
+    private StringPrompt(String string) {
+        this.string = Objects.requireNonNull(string);
+    }
+
+    @Override
+    public String asString() { return string; }
+
+    @Override
+    public StringPrompt append(String text) {
+        return StringPrompt.from(string + text);
+    }
+
+    @Override
+    public StringPrompt append(Completion completion) {
+        return append(completion.text());
+    }
+
+    @Override
+    public String toString() {
+        return string;
+    }
+
+    public static StringPrompt from(String string) {
+        return new StringPrompt(string);
+    }
+
+}
diff --git a/vespajlib/src/main/java/ai/vespa/llm/completion/package-info.java b/vespajlib/src/main/java/ai/vespa/llm/completion/package-info.java
new file mode 100644
index 00000000000..79898c694ca
--- /dev/null
+++ b/vespajlib/src/main/java/ai/vespa/llm/completion/package-info.java
@@ -0,0 +1,11 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+@ExportPackage
+@PublicApi
+package ai.vespa.llm.completion;
+
+import com.yahoo.api.annotations.PublicApi;
+import com.yahoo.osgi.annotation.ExportPackage;
+
+/**
+ * Classes for generating text completions with language models.
+ */
\ No newline at end of file
diff --git a/vespajlib/src/main/java/ai/vespa/llm/package-info.java b/vespajlib/src/main/java/ai/vespa/llm/package-info.java
new file mode 100644
index 00000000000..04fc24c51ee
--- /dev/null
+++ b/vespajlib/src/main/java/ai/vespa/llm/package-info.java
@@ -0,0 +1,11 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+@ExportPackage
+@PublicApi
+package ai.vespa.llm;
+
+import com.yahoo.api.annotations.PublicApi;
+import com.yahoo.osgi.annotation.ExportPackage;
+
+/**
+ * API for working with large language models.
+ */
\ No newline at end of file
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); }
+
+    }
+
+}
diff --git a/vespajlib/src/main/java/ai/vespa/llm/test/package-info.java b/vespajlib/src/main/java/ai/vespa/llm/test/package-info.java
new file mode 100644
index 00000000000..0d51815fd6d
--- /dev/null
+++ b/vespajlib/src/main/java/ai/vespa/llm/test/package-info.java
@@ -0,0 +1,11 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+@ExportPackage
+@PublicApi
+package ai.vespa.llm.test;
+
+/**
+ * Tools for writing tests when working with large language models.
+ */
+
+import com.yahoo.api.annotations.PublicApi;
+import com.yahoo.osgi.annotation.ExportPackage;
\ No newline at end of file
diff --git a/vespajlib/src/test/java/ai/vespa/llm/completion/CompletionTest.java b/vespajlib/src/test/java/ai/vespa/llm/completion/CompletionTest.java
new file mode 100644
index 00000000000..1c794c64d1a
--- /dev/null
+++ b/vespajlib/src/test/java/ai/vespa/llm/completion/CompletionTest.java
@@ -0,0 +1,40 @@
+package ai.vespa.llm.completion;
+
+import ai.vespa.llm.completion.Completion;
+import ai.vespa.llm.completion.Prompt;
+import ai.vespa.llm.completion.StringPrompt;
+import ai.vespa.llm.test.MockLanguageModel;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+import java.util.function.Function;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+/**
+ * Tests completion with a mock completer.
+ *
+ * @author bratseth
+ */
+public class CompletionTest {
+
+    @Test
+    public void testCompletion() {
+        Function<Prompt, List<Completion>> completer = in ->
+            switch (in.asString()) {
+                case "Complete this: " -> List.of(Completion.from("The completion"));
+                default -> throw new RuntimeException("Cannot complete '" + in + "'");
+        };
+        var llm = new MockLanguageModel.Builder().completer(completer).build();
+
+        String input = "Complete this: ";
+        StringPrompt prompt = StringPrompt.from(input);
+        for (int i = 0; i < 10; i++) {
+            var completion = llm.complete(prompt).get(0);
+            prompt = prompt.append(completion);
+            if (completion.finishReason() == Completion.FinishReason.stop) break;
+        }
+        assertEquals("Complete this: The completion", prompt.asString());
+    }
+
+}