migrate vespa to jackson + tests and refactorings (#2056)
## Issue Closes #1681 ## Change 1. Raise the baseline from jdk8 to jdk17 (because the newest vespa-feed-client is no longer support jdk8). 2. Migrate from Gson to Jackson. 3. Get rid of Lombok 4. Implement VespaEmbeddingStoreIT 5. <s>VespaEmbeddingStoreCloudIT</s> will be done in scope of other ticket 6. Support new API search() 7. Support logRequests & logResponses 8. Make parameters keyPath & certPath optional 9. Support removeAll() ## General checklist - [x] There are no breaking changes - [x] I have added unit and integration tests for my change - [x] I have manually run all the unit and integration tests in the module I have added/changed, and they are all green - [x] I have manually run all the unit and integration tests in the [core](https://github.com/langchain4j/langchain4j/tree/main/langchain4j-core) and [main](https://github.com/langchain4j/langchain4j/tree/main/langchain4j) modules, and they are all green - [x] <s>I have added/updated the [documentation](https://github.com/langchain4j/langchain4j/tree/main/docs/docs)</s> will be done in scope of other ticket - [x] I have added an example in the [examples repo](https://github.com/langchain4j/langchain4j-examples) (only for "big" features) will be updated in scope of other ticket - [x] I have added/updated [Spring Boot starter(s)](https://github.com/langchain4j/langchain4j-spring) (if applicable) ## Checklist for changing existing embedding store integration <!-- Please double-check the following points and mark them like this: [X] --> - [x] I have manually verified that the `{NameOfIntegration}EmbeddingStore` works correctly with the data persisted using the latest released version of LangChain4j
This commit is contained in:
Alexey Titov
GitHub
parent
5194e05a02
commit
82ca87bfeb
|
|
@ -1,20 +1,19 @@
|
|||
package dev.langchain4j.store.embedding;
|
||||
|
||||
import dev.langchain4j.data.embedding.Embedding;
|
||||
import dev.langchain4j.data.segment.TextSegment;
|
||||
import dev.langchain4j.model.embedding.EmbeddingModel;
|
||||
import org.assertj.core.data.Percentage;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import static dev.langchain4j.internal.Utils.randomUUID;
|
||||
import static dev.langchain4j.store.embedding.TestUtils.awaitUntilAsserted;
|
||||
import static java.util.Arrays.asList;
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.data.Percentage.withPercentage;
|
||||
|
||||
import dev.langchain4j.data.embedding.Embedding;
|
||||
import dev.langchain4j.data.segment.TextSegment;
|
||||
import dev.langchain4j.model.embedding.EmbeddingModel;
|
||||
import java.util.List;
|
||||
import org.assertj.core.data.Percentage;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
public abstract class EmbeddingStoreWithoutMetadataIT {
|
||||
|
||||
protected abstract EmbeddingStore<TextSegment> embeddingStore();
|
||||
|
|
@ -28,11 +27,9 @@ public abstract class EmbeddingStoreWithoutMetadataIT {
|
|||
ensureStoreIsEmpty();
|
||||
}
|
||||
|
||||
protected void ensureStoreIsReady() {
|
||||
}
|
||||
protected void ensureStoreIsReady() {}
|
||||
|
||||
protected void clearStore() {
|
||||
}
|
||||
protected void clearStore() {}
|
||||
|
||||
protected void ensureStoreIsEmpty() {
|
||||
assertThat(getAllEmbeddings()).isEmpty();
|
||||
|
|
@ -40,7 +37,6 @@ public abstract class EmbeddingStoreWithoutMetadataIT {
|
|||
|
||||
@Test
|
||||
void should_add_embedding() {
|
||||
|
||||
// given
|
||||
Embedding embedding = embeddingModel().embed("hello").content();
|
||||
String id = embeddingStore().add(embedding);
|
||||
|
|
@ -62,15 +58,17 @@ public abstract class EmbeddingStoreWithoutMetadataIT {
|
|||
assertThat(match.embedded()).isNull();
|
||||
|
||||
// new API
|
||||
assertThat(embeddingStore().search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(embedding)
|
||||
.maxResults(10)
|
||||
.build()).matches()).isEqualTo(relevant);
|
||||
assertThat(embeddingStore()
|
||||
.search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(embedding)
|
||||
.maxResults(10)
|
||||
.build())
|
||||
.matches())
|
||||
.isEqualTo(relevant);
|
||||
}
|
||||
|
||||
@Test
|
||||
void should_add_embedding_with_id() {
|
||||
|
||||
// given
|
||||
String id = randomUUID();
|
||||
Embedding embedding = embeddingModel().embed("hello").content();
|
||||
|
|
@ -92,15 +90,17 @@ public abstract class EmbeddingStoreWithoutMetadataIT {
|
|||
assertThat(match.embedded()).isNull();
|
||||
|
||||
// new API
|
||||
assertThat(embeddingStore().search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(embedding)
|
||||
.maxResults(10)
|
||||
.build()).matches()).isEqualTo(relevant);
|
||||
assertThat(embeddingStore()
|
||||
.search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(embedding)
|
||||
.maxResults(10)
|
||||
.build())
|
||||
.matches())
|
||||
.isEqualTo(relevant);
|
||||
}
|
||||
|
||||
@Test
|
||||
void should_add_embedding_with_segment() {
|
||||
|
||||
// given
|
||||
TextSegment segment = TextSegment.from("hello");
|
||||
Embedding embedding = embeddingModel().embed(segment.text()).content();
|
||||
|
|
@ -123,15 +123,17 @@ public abstract class EmbeddingStoreWithoutMetadataIT {
|
|||
assertThat(match.embedded()).isEqualTo(segment);
|
||||
|
||||
// new API
|
||||
assertThat(embeddingStore().search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(embedding)
|
||||
.maxResults(10)
|
||||
.build()).matches()).isEqualTo(relevant);
|
||||
assertThat(embeddingStore()
|
||||
.search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(embedding)
|
||||
.maxResults(10)
|
||||
.build())
|
||||
.matches())
|
||||
.isEqualTo(relevant);
|
||||
}
|
||||
|
||||
@Test
|
||||
void should_add_multiple_embeddings() {
|
||||
|
||||
// given
|
||||
Embedding firstEmbedding = embeddingModel().embed("hello").content();
|
||||
Embedding secondEmbedding = embeddingModel().embed("hi").content();
|
||||
|
|
@ -158,10 +160,10 @@ public abstract class EmbeddingStoreWithoutMetadataIT {
|
|||
assertThat(firstMatch.embedded()).isNull();
|
||||
|
||||
EmbeddingMatch<TextSegment> secondMatch = relevant.get(1);
|
||||
assertThat(secondMatch.score()).isCloseTo(
|
||||
RelevanceScore.fromCosineSimilarity(CosineSimilarity.between(firstEmbedding, secondEmbedding)),
|
||||
percentage()
|
||||
);
|
||||
assertThat(secondMatch.score())
|
||||
.isCloseTo(
|
||||
RelevanceScore.fromCosineSimilarity(CosineSimilarity.between(firstEmbedding, secondEmbedding)),
|
||||
percentage());
|
||||
assertThat(secondMatch.embeddingId()).isEqualTo(ids.get(1));
|
||||
if (assertEmbedding()) {
|
||||
assertThat(CosineSimilarity.between(secondMatch.embedding(), secondEmbedding))
|
||||
|
|
@ -170,15 +172,17 @@ public abstract class EmbeddingStoreWithoutMetadataIT {
|
|||
assertThat(secondMatch.embedded()).isNull();
|
||||
|
||||
// new API
|
||||
assertThat(embeddingStore().search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(firstEmbedding)
|
||||
.maxResults(10)
|
||||
.build()).matches()).isEqualTo(relevant);
|
||||
assertThat(embeddingStore()
|
||||
.search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(firstEmbedding)
|
||||
.maxResults(10)
|
||||
.build())
|
||||
.matches())
|
||||
.isEqualTo(relevant);
|
||||
}
|
||||
|
||||
@Test
|
||||
void should_add_multiple_embeddings_with_segments() {
|
||||
|
||||
// given
|
||||
TextSegment firstSegment = TextSegment.from("hello");
|
||||
Embedding firstEmbedding = embeddingModel().embed(firstSegment.text()).content();
|
||||
|
|
@ -186,10 +190,8 @@ public abstract class EmbeddingStoreWithoutMetadataIT {
|
|||
TextSegment secondSegment = TextSegment.from("hi");
|
||||
Embedding secondEmbedding = embeddingModel().embed(secondSegment.text()).content();
|
||||
|
||||
List<String> ids = embeddingStore().addAll(
|
||||
asList(firstEmbedding, secondEmbedding),
|
||||
asList(firstSegment, secondSegment)
|
||||
);
|
||||
List<String> ids =
|
||||
embeddingStore().addAll(asList(firstEmbedding, secondEmbedding), asList(firstSegment, secondSegment));
|
||||
|
||||
awaitUntilAsserted(() -> assertThat(getAllEmbeddings()).hasSize(2));
|
||||
|
||||
|
|
@ -212,10 +214,10 @@ public abstract class EmbeddingStoreWithoutMetadataIT {
|
|||
assertThat(firstMatch.embedded()).isEqualTo(firstSegment);
|
||||
|
||||
EmbeddingMatch<TextSegment> secondMatch = relevant.get(1);
|
||||
assertThat(secondMatch.score()).isCloseTo(
|
||||
RelevanceScore.fromCosineSimilarity(CosineSimilarity.between(firstEmbedding, secondEmbedding)),
|
||||
percentage()
|
||||
);
|
||||
assertThat(secondMatch.score())
|
||||
.isCloseTo(
|
||||
RelevanceScore.fromCosineSimilarity(CosineSimilarity.between(firstEmbedding, secondEmbedding)),
|
||||
percentage());
|
||||
assertThat(secondMatch.embeddingId()).isEqualTo(ids.get(1));
|
||||
if (assertEmbedding()) {
|
||||
assertThat(CosineSimilarity.between(secondMatch.embedding(), secondEmbedding))
|
||||
|
|
@ -224,10 +226,13 @@ public abstract class EmbeddingStoreWithoutMetadataIT {
|
|||
assertThat(secondMatch.embedded()).isEqualTo(secondSegment);
|
||||
|
||||
// new API
|
||||
assertThat(embeddingStore().search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(firstEmbedding)
|
||||
.maxResults(10)
|
||||
.build()).matches()).isEqualTo(relevant);
|
||||
assertThat(embeddingStore()
|
||||
.search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(firstEmbedding)
|
||||
.maxResults(10)
|
||||
.build())
|
||||
.matches())
|
||||
.isEqualTo(relevant);
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -243,11 +248,8 @@ public abstract class EmbeddingStoreWithoutMetadataIT {
|
|||
TextSegment secondSegment = TextSegment.from("hi");
|
||||
Embedding secondEmbedding = embeddingModel().embed(secondSegment.text()).content();
|
||||
|
||||
embeddingStore().addAll(
|
||||
asList(id1, id2),
|
||||
asList(firstEmbedding, secondEmbedding),
|
||||
asList(firstSegment, secondSegment)
|
||||
);
|
||||
embeddingStore()
|
||||
.addAll(asList(id1, id2), asList(firstEmbedding, secondEmbedding), asList(firstSegment, secondSegment));
|
||||
|
||||
awaitUntilAsserted(() -> assertThat(getAllEmbeddings()).hasSize(2));
|
||||
|
||||
|
|
@ -270,10 +272,10 @@ public abstract class EmbeddingStoreWithoutMetadataIT {
|
|||
assertThat(firstMatch.embedded()).isEqualTo(firstSegment);
|
||||
|
||||
EmbeddingMatch<TextSegment> secondMatch = relevant.get(1);
|
||||
assertThat(secondMatch.score()).isCloseTo(
|
||||
RelevanceScore.fromCosineSimilarity(CosineSimilarity.between(firstEmbedding, secondEmbedding)),
|
||||
percentage()
|
||||
);
|
||||
assertThat(secondMatch.score())
|
||||
.isCloseTo(
|
||||
RelevanceScore.fromCosineSimilarity(CosineSimilarity.between(firstEmbedding, secondEmbedding)),
|
||||
percentage());
|
||||
assertThat(secondMatch.embeddingId()).isEqualTo(id2);
|
||||
if (assertEmbedding()) {
|
||||
assertThat(CosineSimilarity.between(secondMatch.embedding(), secondEmbedding))
|
||||
|
|
@ -282,15 +284,17 @@ public abstract class EmbeddingStoreWithoutMetadataIT {
|
|||
assertThat(secondMatch.embedded()).isEqualTo(secondSegment);
|
||||
|
||||
// new API
|
||||
assertThat(embeddingStore().search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(firstEmbedding)
|
||||
.maxResults(10)
|
||||
.build()).matches()).isEqualTo(relevant);
|
||||
assertThat(embeddingStore()
|
||||
.search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(firstEmbedding)
|
||||
.maxResults(10)
|
||||
.build())
|
||||
.matches())
|
||||
.isEqualTo(relevant);
|
||||
}
|
||||
|
||||
@Test
|
||||
void should_find_with_min_score() {
|
||||
|
||||
// given
|
||||
String firstId = randomUUID();
|
||||
Embedding firstEmbedding = embeddingModel().embed("hello").content();
|
||||
|
|
@ -311,24 +315,24 @@ public abstract class EmbeddingStoreWithoutMetadataIT {
|
|||
assertThat(firstMatch.score()).isCloseTo(1, percentage());
|
||||
assertThat(firstMatch.embeddingId()).isEqualTo(firstId);
|
||||
EmbeddingMatch<TextSegment> secondMatch = relevant.get(1);
|
||||
assertThat(secondMatch.score()).isCloseTo(
|
||||
RelevanceScore.fromCosineSimilarity(CosineSimilarity.between(firstEmbedding, secondEmbedding)),
|
||||
percentage()
|
||||
);
|
||||
assertThat(secondMatch.score())
|
||||
.isCloseTo(
|
||||
RelevanceScore.fromCosineSimilarity(CosineSimilarity.between(firstEmbedding, secondEmbedding)),
|
||||
percentage());
|
||||
assertThat(secondMatch.embeddingId()).isEqualTo(secondId);
|
||||
|
||||
// new API
|
||||
assertThat(embeddingStore().search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(firstEmbedding)
|
||||
.maxResults(10)
|
||||
.build()).matches()).isEqualTo(relevant);
|
||||
assertThat(embeddingStore()
|
||||
.search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(firstEmbedding)
|
||||
.maxResults(10)
|
||||
.build())
|
||||
.matches())
|
||||
.isEqualTo(relevant);
|
||||
|
||||
// when
|
||||
List<EmbeddingMatch<TextSegment>> relevant2 = embeddingStore().findRelevant(
|
||||
firstEmbedding,
|
||||
10,
|
||||
secondMatch.score() - 0.01
|
||||
);
|
||||
List<EmbeddingMatch<TextSegment>> relevant2 =
|
||||
embeddingStore().findRelevant(firstEmbedding, 10, secondMatch.score() - 0.01);
|
||||
|
||||
// then
|
||||
assertThat(relevant2).hasSize(2);
|
||||
|
|
@ -336,18 +340,18 @@ public abstract class EmbeddingStoreWithoutMetadataIT {
|
|||
assertThat(relevant2.get(1).embeddingId()).isEqualTo(secondId);
|
||||
|
||||
// new API
|
||||
assertThat(embeddingStore().search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(firstEmbedding)
|
||||
.maxResults(10)
|
||||
.minScore(secondMatch.score() - 0.01)
|
||||
.build()).matches()).isEqualTo(relevant2);
|
||||
assertThat(embeddingStore()
|
||||
.search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(firstEmbedding)
|
||||
.maxResults(10)
|
||||
.minScore(secondMatch.score() - 0.01)
|
||||
.build())
|
||||
.matches())
|
||||
.isEqualTo(relevant2);
|
||||
|
||||
// when
|
||||
List<EmbeddingMatch<TextSegment>> relevant3 = embeddingStore().findRelevant(
|
||||
firstEmbedding,
|
||||
10,
|
||||
secondMatch.score()
|
||||
);
|
||||
List<EmbeddingMatch<TextSegment>> relevant3 =
|
||||
embeddingStore().findRelevant(firstEmbedding, 10, secondMatch.score());
|
||||
|
||||
// then
|
||||
assertThat(relevant3).hasSize(2);
|
||||
|
|
@ -355,34 +359,36 @@ public abstract class EmbeddingStoreWithoutMetadataIT {
|
|||
assertThat(relevant3.get(1).embeddingId()).isEqualTo(secondId);
|
||||
|
||||
// new API
|
||||
assertThat(embeddingStore().search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(firstEmbedding)
|
||||
.maxResults(10)
|
||||
.minScore(secondMatch.score())
|
||||
.build()).matches()).isEqualTo(relevant3);
|
||||
assertThat(embeddingStore()
|
||||
.search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(firstEmbedding)
|
||||
.maxResults(10)
|
||||
.minScore(secondMatch.score())
|
||||
.build())
|
||||
.matches())
|
||||
.isEqualTo(relevant3);
|
||||
|
||||
// when
|
||||
List<EmbeddingMatch<TextSegment>> relevant4 = embeddingStore().findRelevant(
|
||||
firstEmbedding,
|
||||
10,
|
||||
secondMatch.score() + 0.01
|
||||
);
|
||||
List<EmbeddingMatch<TextSegment>> relevant4 =
|
||||
embeddingStore().findRelevant(firstEmbedding, 10, secondMatch.score() + 0.01);
|
||||
|
||||
// then
|
||||
assertThat(relevant4).hasSize(1);
|
||||
assertThat(relevant4.get(0).embeddingId()).isEqualTo(firstId);
|
||||
|
||||
// new API
|
||||
assertThat(embeddingStore().search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(firstEmbedding)
|
||||
.maxResults(10)
|
||||
.minScore(secondMatch.score() + 0.01)
|
||||
.build()).matches()).isEqualTo(relevant4);
|
||||
assertThat(embeddingStore()
|
||||
.search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(firstEmbedding)
|
||||
.maxResults(10)
|
||||
.minScore(secondMatch.score() + 0.01)
|
||||
.build())
|
||||
.matches())
|
||||
.isEqualTo(relevant4);
|
||||
}
|
||||
|
||||
@Test
|
||||
void should_return_correct_score() {
|
||||
|
||||
// given
|
||||
Embedding embedding = embeddingModel().embed("hello").content();
|
||||
|
||||
|
|
@ -399,20 +405,22 @@ public abstract class EmbeddingStoreWithoutMetadataIT {
|
|||
// then
|
||||
assertThat(relevant).hasSize(1);
|
||||
EmbeddingMatch<TextSegment> match = relevant.get(0);
|
||||
assertThat(match.score()).isCloseTo(
|
||||
RelevanceScore.fromCosineSimilarity(CosineSimilarity.between(embedding, referenceEmbedding)),
|
||||
percentage()
|
||||
);
|
||||
assertThat(match.score())
|
||||
.isCloseTo(
|
||||
RelevanceScore.fromCosineSimilarity(CosineSimilarity.between(embedding, referenceEmbedding)),
|
||||
percentage());
|
||||
|
||||
// new API
|
||||
assertThat(embeddingStore().search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(referenceEmbedding)
|
||||
.maxResults(1)
|
||||
.build()).matches()).isEqualTo(relevant);
|
||||
assertThat(embeddingStore()
|
||||
.search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(referenceEmbedding)
|
||||
.maxResults(1)
|
||||
.build())
|
||||
.matches())
|
||||
.isEqualTo(relevant);
|
||||
}
|
||||
|
||||
protected List<EmbeddingMatch<TextSegment>> getAllEmbeddings() {
|
||||
|
||||
EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(embeddingModel().embed("test").content())
|
||||
.maxResults(1000)
|
||||
|
|
|
|||
|
|
@ -28,8 +28,7 @@
|
|||
</licenses>
|
||||
|
||||
<properties>
|
||||
<vespa.version>8.190.2</vespa.version>
|
||||
<!-- the latest Java 8 version -->
|
||||
<vespa.version>8.458.13</vespa.version>
|
||||
</properties>
|
||||
|
||||
<dependencies>
|
||||
|
|
@ -59,7 +58,17 @@
|
|||
|
||||
<dependency>
|
||||
<groupId>com.squareup.retrofit2</groupId>
|
||||
<artifactId>converter-gson</artifactId>
|
||||
<artifactId>converter-jackson</artifactId>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>com.fasterxml.jackson.core</groupId>
|
||||
<artifactId>jackson-databind</artifactId>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.slf4j</groupId>
|
||||
<artifactId>slf4j-api</artifactId>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
|
|
@ -67,12 +76,6 @@
|
|||
<artifactId>okhttp</artifactId>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.projectlombok</groupId>
|
||||
<artifactId>lombok</artifactId>
|
||||
<scope>provided</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>dev.langchain4j</groupId>
|
||||
<artifactId>langchain4j-core</artifactId>
|
||||
|
|
@ -102,6 +105,41 @@
|
|||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.assertj</groupId>
|
||||
<artifactId>assertj-core</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.tinylog</groupId>
|
||||
<artifactId>tinylog-impl</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.tinylog</groupId>
|
||||
<artifactId>slf4j-tinylog</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.awaitility</groupId>
|
||||
<artifactId>awaitility</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.testcontainers</groupId>
|
||||
<artifactId>testcontainers</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.testcontainers</groupId>
|
||||
<artifactId>junit-jupiter</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
|
||||
</project>
|
||||
|
|
|
|||
|
|
@ -0,0 +1,6 @@
|
|||
package dev.langchain4j.store.embedding.vespa;
|
||||
|
||||
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
|
||||
|
||||
@JsonIgnoreProperties(ignoreUnknown = true)
|
||||
public record DeleteResponse(String pathId, Long documentCount) {}
|
||||
|
|
@ -1,29 +1,10 @@
|
|||
package dev.langchain4j.store.embedding.vespa;
|
||||
|
||||
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
|
||||
import java.util.List;
|
||||
|
||||
class QueryResponse {
|
||||
|
||||
private RootNode root;
|
||||
|
||||
public RootNode getRoot() {
|
||||
return root;
|
||||
}
|
||||
|
||||
public void setRoot(RootNode root) {
|
||||
this.root = root;
|
||||
}
|
||||
|
||||
public static class RootNode {
|
||||
|
||||
private List<Record> children;
|
||||
|
||||
public List<Record> getChildren() {
|
||||
return children;
|
||||
}
|
||||
|
||||
public void setChildren(List<Record> children) {
|
||||
this.children = children;
|
||||
}
|
||||
}
|
||||
@JsonIgnoreProperties(ignoreUnknown = true)
|
||||
record QueryResponse(RootNode root) {
|
||||
@JsonIgnoreProperties(ignoreUnknown = true)
|
||||
record RootNode(List<Record> children) {}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,97 +1,21 @@
|
|||
package dev.langchain4j.store.embedding.vespa;
|
||||
|
||||
import com.google.gson.annotations.SerializedName;
|
||||
import static com.fasterxml.jackson.annotation.JsonInclude.Include.NON_NULL;
|
||||
import static com.fasterxml.jackson.databind.PropertyNamingStrategies.SnakeCaseStrategy;
|
||||
|
||||
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
|
||||
import com.fasterxml.jackson.annotation.JsonInclude;
|
||||
import com.fasterxml.jackson.databind.annotation.JsonNaming;
|
||||
import java.util.List;
|
||||
|
||||
class Record {
|
||||
|
||||
private String id;
|
||||
private Double relevance;
|
||||
private Fields fields;
|
||||
|
||||
public Record(String id, String textSegment, List<Float> vector) {
|
||||
this.id = id;
|
||||
this.fields = new Fields(textSegment, vector);
|
||||
}
|
||||
|
||||
public String getId() {
|
||||
return id;
|
||||
}
|
||||
|
||||
public void setId(String id) {
|
||||
this.id = id;
|
||||
}
|
||||
|
||||
public double getRelevance() {
|
||||
return relevance;
|
||||
}
|
||||
|
||||
public void setRelevance(double relevance) {
|
||||
this.relevance = relevance;
|
||||
}
|
||||
|
||||
public Fields getFields() {
|
||||
return fields;
|
||||
}
|
||||
|
||||
public void setFields(Fields fields) {
|
||||
this.fields = fields;
|
||||
}
|
||||
|
||||
public static class Fields {
|
||||
|
||||
@SerializedName("documentid")
|
||||
private String documentId;
|
||||
|
||||
@SerializedName("text_segment")
|
||||
private String textSegment;
|
||||
|
||||
private Vector vector;
|
||||
|
||||
public Fields(String textSegment, List<Float> vector) {
|
||||
this.textSegment = textSegment;
|
||||
this.vector = new Vector(vector);
|
||||
@JsonInclude(NON_NULL)
|
||||
@JsonIgnoreProperties(ignoreUnknown = true)
|
||||
public record Record(String id, Double relevance, Fields fields) {
|
||||
@JsonIgnoreProperties(ignoreUnknown = true)
|
||||
@JsonInclude(NON_NULL)
|
||||
@JsonNaming(SnakeCaseStrategy.class)
|
||||
public record Fields(String documentid, String textSegment, Vector vector) {
|
||||
@JsonIgnoreProperties(ignoreUnknown = true)
|
||||
public record Vector(List<Float> values) {}
|
||||
}
|
||||
|
||||
public String getDocumentId() {
|
||||
return documentId;
|
||||
}
|
||||
|
||||
public void setDocumentId(String documentId) {
|
||||
this.documentId = documentId;
|
||||
}
|
||||
|
||||
public String getTextSegment() {
|
||||
return textSegment;
|
||||
}
|
||||
|
||||
public void setTextSegment(String textSegment) {
|
||||
this.textSegment = textSegment;
|
||||
}
|
||||
|
||||
public Vector getVector() {
|
||||
return vector;
|
||||
}
|
||||
|
||||
public void setVector(Vector vector) {
|
||||
this.vector = vector;
|
||||
}
|
||||
|
||||
public static class Vector {
|
||||
|
||||
private List<Float> values;
|
||||
|
||||
public Vector(List<Float> values) {
|
||||
this.values = values;
|
||||
}
|
||||
|
||||
public List<Float> getValues() {
|
||||
return values;
|
||||
}
|
||||
|
||||
public void setValues(List<Float> values) {
|
||||
this.values = values;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,16 @@
|
|||
package dev.langchain4j.store.embedding.vespa;
|
||||
|
||||
import retrofit2.Call;
|
||||
import retrofit2.http.DELETE;
|
||||
import retrofit2.http.GET;
|
||||
import retrofit2.http.Path;
|
||||
import retrofit2.http.Query;
|
||||
|
||||
interface VespaApi {
|
||||
@GET("search/{query}")
|
||||
Call<QueryResponse> search(@Path(value = "query", encoded = true) String query);
|
||||
|
||||
@DELETE("document/v1/{ns}/{docType}/docid?selection=true")
|
||||
Call<DeleteResponse> deleteAll(
|
||||
@Path("ns") String namespace, @Path("docType") String documentType, @Query("cluster") String clusterName);
|
||||
}
|
||||
|
|
@ -0,0 +1,162 @@
|
|||
// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
|
||||
package dev.langchain4j.store.embedding.vespa;
|
||||
|
||||
import dev.langchain4j.internal.Utils;
|
||||
import java.io.IOException;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.security.GeneralSecurityException;
|
||||
import java.security.KeyFactory;
|
||||
import java.security.KeyStore;
|
||||
import java.security.PrivateKey;
|
||||
import java.security.cert.Certificate;
|
||||
import java.security.cert.X509Certificate;
|
||||
import java.security.spec.PKCS8EncodedKeySpec;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import javax.net.ssl.KeyManager;
|
||||
import javax.net.ssl.KeyManagerFactory;
|
||||
import javax.net.ssl.SSLContext;
|
||||
import javax.net.ssl.TrustManagerFactory;
|
||||
import javax.net.ssl.X509TrustManager;
|
||||
import okhttp3.HttpUrl;
|
||||
import okhttp3.OkHttpClient;
|
||||
import okhttp3.Request;
|
||||
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
|
||||
import org.bouncycastle.asn1.pkcs.PKCSObjectIdentifiers;
|
||||
import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
|
||||
import org.bouncycastle.asn1.x9.X9ObjectIdentifiers;
|
||||
import org.bouncycastle.cert.X509CertificateHolder;
|
||||
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
|
||||
import org.bouncycastle.jce.provider.BouncyCastleProvider;
|
||||
import org.bouncycastle.openssl.PEMKeyPair;
|
||||
import org.bouncycastle.openssl.PEMParser;
|
||||
import retrofit2.Retrofit;
|
||||
import retrofit2.converter.jackson.JacksonConverterFactory;
|
||||
|
||||
/**
|
||||
* This Workaround is needed because of <a href="https://github.com/vespa-engine/vespa/issues/28026">this request</a>.
|
||||
* It will be redundant as soon as vespa-client is implemented. This class is copied from <code>vespa-feed-client</code>.
|
||||
* BouncyCastle integration for creating a {@link SSLContext} instance from PEM encoded material
|
||||
*/
|
||||
class VespaClient {
|
||||
|
||||
static final BouncyCastleProvider bcProvider = new BouncyCastleProvider();
|
||||
|
||||
public static VespaApi createInstance(
|
||||
String baseUrl, Path certificate, Path privateKey, boolean logRequests, boolean logResponses) {
|
||||
try {
|
||||
OkHttpClient.Builder builder = new OkHttpClient.Builder().addInterceptor(chain -> {
|
||||
// trick to format the query URL exactly how Vespa expects it (search/?query),
|
||||
// see https://docs.vespa.ai/en/reference/query-language-reference.html
|
||||
Request request = chain.request();
|
||||
if (request.url().url().getPath().startsWith("/search/")) {
|
||||
HttpUrl url = request.url()
|
||||
.newBuilder()
|
||||
.removePathSegment(1)
|
||||
.addPathSegment("")
|
||||
.encodedQuery(request.url().encodedPathSegments().get(1))
|
||||
.build();
|
||||
request = request.newBuilder().url(url).build();
|
||||
}
|
||||
return chain.proceed(request);
|
||||
});
|
||||
|
||||
addSsl(certificate, privateKey, builder);
|
||||
|
||||
if (logRequests) {
|
||||
builder.addInterceptor(new VespaRequestLoggingInterceptor());
|
||||
}
|
||||
if (logResponses) {
|
||||
builder.addInterceptor(new VespaResponseLoggingInterceptor());
|
||||
}
|
||||
|
||||
OkHttpClient client = builder.build();
|
||||
|
||||
Retrofit retrofit = new Retrofit.Builder()
|
||||
.baseUrl(Utils.ensureTrailingForwardSlash(baseUrl))
|
||||
.client(client)
|
||||
.addConverterFactory(JacksonConverterFactory.create())
|
||||
.build();
|
||||
|
||||
return retrofit.create(VespaApi.class);
|
||||
} catch (Exception e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
private static void addSsl(Path certificate, Path privateKey, OkHttpClient.Builder builder)
|
||||
throws IOException, GeneralSecurityException {
|
||||
if (certificate != null && privateKey != null) {
|
||||
KeyStore keystore = KeyStore.getInstance("PKCS12");
|
||||
keystore.load(null);
|
||||
keystore.setKeyEntry("cert", privateKey(privateKey), new char[0], certificates(certificate));
|
||||
// Protocol version must be equal to TlsContext.SSL_CONTEXT_VERSION or higher
|
||||
SSLContext sslContext = SSLContext.getInstance("TLSv1.3");
|
||||
sslContext.init(createKeyManagers(keystore), null, /*Default secure random algorithm*/ null);
|
||||
|
||||
TrustManagerFactory trustManagerFactory =
|
||||
TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
|
||||
trustManagerFactory.init(keystore);
|
||||
|
||||
builder.sslSocketFactory(sslContext.getSocketFactory(), (X509TrustManager)
|
||||
trustManagerFactory.getTrustManagers()[0]);
|
||||
}
|
||||
}
|
||||
|
||||
private static KeyManager[] createKeyManagers(KeyStore keystore) throws GeneralSecurityException {
|
||||
KeyManagerFactory kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
|
||||
kmf.init(keystore, new char[0]);
|
||||
return kmf.getKeyManagers();
|
||||
}
|
||||
|
||||
private static Certificate[] certificates(Path file) throws IOException, GeneralSecurityException {
|
||||
try (PEMParser parser = new PEMParser(Files.newBufferedReader(file))) {
|
||||
List<X509Certificate> result = new ArrayList<>();
|
||||
Object pemObject;
|
||||
while ((pemObject = parser.readObject()) != null) {
|
||||
result.add(toX509Certificate(pemObject));
|
||||
}
|
||||
if (result.isEmpty()) throw new IOException("File contains no PEM encoded certificates: " + file);
|
||||
return result.toArray(new Certificate[0]);
|
||||
}
|
||||
}
|
||||
|
||||
private static PrivateKey privateKey(Path file) throws IOException, GeneralSecurityException {
|
||||
try (PEMParser parser = new PEMParser(Files.newBufferedReader(file))) {
|
||||
Object pemObject;
|
||||
while ((pemObject = parser.readObject()) != null) {
|
||||
if (pemObject instanceof PrivateKeyInfo) {
|
||||
PrivateKeyInfo keyInfo = (PrivateKeyInfo) pemObject;
|
||||
PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyInfo.getEncoded());
|
||||
return createKeyFactory(keyInfo).generatePrivate(keySpec);
|
||||
} else if (pemObject instanceof PEMKeyPair) {
|
||||
PEMKeyPair pemKeypair = (PEMKeyPair) pemObject;
|
||||
PrivateKeyInfo keyInfo = pemKeypair.getPrivateKeyInfo();
|
||||
return createKeyFactory(keyInfo).generatePrivate(new PKCS8EncodedKeySpec(keyInfo.getEncoded()));
|
||||
}
|
||||
}
|
||||
throw new IOException("Could not find private key in PEM file");
|
||||
}
|
||||
}
|
||||
|
||||
private static X509Certificate toX509Certificate(Object pemObject) throws IOException, GeneralSecurityException {
|
||||
if (pemObject instanceof X509Certificate) return (X509Certificate) pemObject;
|
||||
if (pemObject instanceof X509CertificateHolder) {
|
||||
return new JcaX509CertificateConverter().setProvider(bcProvider).getCertificate((X509CertificateHolder)
|
||||
pemObject);
|
||||
}
|
||||
throw new IOException("Invalid type of PEM object: " + pemObject);
|
||||
}
|
||||
|
||||
private static KeyFactory createKeyFactory(PrivateKeyInfo info) throws IOException, GeneralSecurityException {
|
||||
ASN1ObjectIdentifier algorithm = info.getPrivateKeyAlgorithm().getAlgorithm();
|
||||
if (X9ObjectIdentifiers.id_ecPublicKey.equals(algorithm)) {
|
||||
return KeyFactory.getInstance("EC", bcProvider);
|
||||
} else if (PKCSObjectIdentifiers.rsaEncryption.equals(algorithm)) {
|
||||
return KeyFactory.getInstance("RSA", bcProvider);
|
||||
} else {
|
||||
throw new IOException("Unknown key algorithm: " + algorithm);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,22 +1,30 @@
|
|||
package dev.langchain4j.store.embedding.vespa;
|
||||
|
||||
import static dev.langchain4j.internal.Utils.generateUUIDFrom;
|
||||
import static dev.langchain4j.internal.Utils.getOrDefault;
|
||||
import static dev.langchain4j.internal.Utils.randomUUID;
|
||||
import static dev.langchain4j.store.embedding.vespa.VespaQueryClient.createInstance;
|
||||
import static dev.langchain4j.internal.ValidationUtils.ensureNotNull;
|
||||
import static dev.langchain4j.store.embedding.vespa.Record.Fields.Vector;
|
||||
import static dev.langchain4j.store.embedding.vespa.VespaClient.createInstance;
|
||||
|
||||
import ai.vespa.client.dsl.A;
|
||||
import ai.vespa.client.dsl.Annotation;
|
||||
import ai.vespa.client.dsl.NearestNeighbor;
|
||||
import ai.vespa.client.dsl.Q;
|
||||
import ai.vespa.feed.client.*;
|
||||
import ai.vespa.feed.client.DocumentId;
|
||||
import ai.vespa.feed.client.FeedClientBuilder;
|
||||
import ai.vespa.feed.client.FeedException;
|
||||
import ai.vespa.feed.client.JsonFeeder;
|
||||
import ai.vespa.feed.client.Result;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import dev.langchain4j.data.embedding.Embedding;
|
||||
import dev.langchain4j.data.segment.TextSegment;
|
||||
import dev.langchain4j.internal.Json;
|
||||
import dev.langchain4j.store.embedding.EmbeddingMatch;
|
||||
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
|
||||
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
|
||||
import dev.langchain4j.store.embedding.EmbeddingStore;
|
||||
import dev.langchain4j.store.embedding.vespa.Record.Fields;
|
||||
import java.io.ByteArrayInputStream;
|
||||
import java.io.IOException;
|
||||
import java.lang.reflect.InvocationTargetException;
|
||||
import java.lang.reflect.Method;
|
||||
import java.net.URI;
|
||||
import java.nio.file.Path;
|
||||
import java.nio.file.Paths;
|
||||
|
|
@ -24,9 +32,7 @@ import java.time.Duration;
|
|||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.stream.Collectors;
|
||||
import lombok.Builder;
|
||||
import lombok.SneakyThrows;
|
||||
import okhttp3.ResponseBody;
|
||||
import retrofit2.Response;
|
||||
|
||||
/**
|
||||
|
|
@ -37,240 +43,403 @@ import retrofit2.Response;
|
|||
*/
|
||||
public class VespaEmbeddingStore implements EmbeddingStore<TextSegment> {
|
||||
|
||||
private static final Duration DEFAULT_TIMEOUT = Duration.ofSeconds(5);
|
||||
private static final String DEFAULT_NAMESPACE = "namespace";
|
||||
private static final String DEFAULT_DOCUMENT_TYPE = "langchain4j";
|
||||
private static final boolean DEFAULT_AVOID_DUPS = true;
|
||||
private static final String FIELD_NAME_TEXT_SEGMENT = "text_segment";
|
||||
private static final String FIELD_NAME_VECTOR = "vector";
|
||||
private static final String FIELD_NAME_DOCUMENT_ID = "documentid";
|
||||
private static final String DEFAULT_RANK_PROFILE = "cosine_similarity";
|
||||
private static final int DEFAULT_TARGET_HITS = 10;
|
||||
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
|
||||
|
||||
private final String url;
|
||||
private final Path keyPath;
|
||||
private final Path certPath;
|
||||
private final Duration timeout;
|
||||
private final String namespace;
|
||||
private final String documentType;
|
||||
private final String rankProfile;
|
||||
private final int targetHits;
|
||||
private final boolean avoidDups;
|
||||
private static final Duration DEFAULT_TIMEOUT = Duration.ofSeconds(5);
|
||||
static final String DEFAULT_NAMESPACE = "namespace";
|
||||
static final String DEFAULT_DOCUMENT_TYPE = "langchain4j";
|
||||
private static final String DEFAULT_CLUSTER_NAME = "langchain4j";
|
||||
private static final boolean DEFAULT_AVOID_DUPS = true;
|
||||
private static final String FIELD_NAME_TEXT_SEGMENT = "text_segment";
|
||||
private static final String FIELD_NAME_VECTOR = "vector";
|
||||
private static final String FIELD_NAME_DOCUMENT_ID = "documentid";
|
||||
private static final String DEFAULT_RANK_PROFILE = "langchain4j_relevance_score";
|
||||
private static final int DEFAULT_TARGET_HITS = 10;
|
||||
|
||||
private VespaQueryApi queryApi;
|
||||
private final String url;
|
||||
private final Path keyPath;
|
||||
private final Path certPath;
|
||||
private final Duration timeout;
|
||||
private final String namespace;
|
||||
private final String documentType;
|
||||
private final String clusterName;
|
||||
private final String rankProfile;
|
||||
private final int targetHits;
|
||||
private final boolean avoidDups;
|
||||
private final boolean logRequests;
|
||||
private final boolean logResponses;
|
||||
|
||||
/**
|
||||
* Creates a new VespaEmbeddingStore instance.
|
||||
*
|
||||
* @param url server url, local or cloud one. The latter you can find under Endpoint of your Vespa
|
||||
* application, e.g. https://alexey-heezer.langchain4j.mytenant346.aws-us-east-1c.dev.z.vespa-app.cloud/
|
||||
* @param keyPath local path to the SSL private key file in PEM format. Read
|
||||
* <a href="https://cloud.vespa.ai/en/getting-started-java">docs</a> for details.
|
||||
* @param certPath local path to the SSL certificate file in PEM format. Read
|
||||
* <a href="https://cloud.vespa.ai/en/getting-started-java">docs</a> for details.
|
||||
* @param timeout for Vespa Java client in <code>java.time.Duration</code> format.
|
||||
* @param namespace required for document ID generation, find more details
|
||||
* <a href="https://docs.vespa.ai/en/documents.html#namespace">here</a>.
|
||||
* @param documentType document type, used for document ID generation, find more details
|
||||
* <a href="https://docs.vespa.ai/en/documents.html#namespace">here</a> and data querying
|
||||
* @param rankProfile rank profile from your .sd schema. Provided example schema configures cosine similarity match
|
||||
* @param targetHits sets the number of hits (10 is default) exposed to the real Vespa's first-phase ranking
|
||||
* function per content node, find more details
|
||||
* <a href="https://docs.vespa.ai/en/nearest-neighbor-search.html#querying-using-nearestneighbor-query-operator">here</a>.
|
||||
* @param avoidDups if true (default), then <code>VespaEmbeddingStore</code> will generate a hashed ID based on
|
||||
* provided text segment, which avoids duplicated entries in DB.
|
||||
* If false, then random ID will be generated.
|
||||
*/
|
||||
@Builder
|
||||
public VespaEmbeddingStore(
|
||||
String url,
|
||||
String keyPath,
|
||||
String certPath,
|
||||
Duration timeout,
|
||||
String namespace,
|
||||
String documentType,
|
||||
String rankProfile,
|
||||
Integer targetHits,
|
||||
Boolean avoidDups
|
||||
) {
|
||||
this.url = url;
|
||||
this.keyPath = Paths.get(keyPath);
|
||||
this.certPath = Paths.get(certPath);
|
||||
this.timeout = timeout != null ? timeout : DEFAULT_TIMEOUT;
|
||||
this.namespace = namespace != null ? namespace : DEFAULT_NAMESPACE;
|
||||
this.documentType = documentType != null ? documentType : DEFAULT_DOCUMENT_TYPE;
|
||||
this.rankProfile = rankProfile != null ? rankProfile : DEFAULT_RANK_PROFILE;
|
||||
this.targetHits = targetHits != null ? targetHits : DEFAULT_TARGET_HITS;
|
||||
this.avoidDups = avoidDups != null ? avoidDups : DEFAULT_AVOID_DUPS;
|
||||
}
|
||||
private VespaApi api;
|
||||
|
||||
@Override
|
||||
public String add(Embedding embedding) {
|
||||
return add(null, embedding, null);
|
||||
}
|
||||
/**
|
||||
* Creates a new VespaEmbeddingStore instance.
|
||||
*
|
||||
* @param url server url, local or cloud one. The latter you can find under Endpoint of your Vespa
|
||||
* application, e.g. https://alexey-heezer.langchain4j.mytenant346.aws-us-east-1c.dev.z.vespa-app.cloud/
|
||||
* @param keyPath local path to the SSL private key file in PEM format. Read
|
||||
* <a href="https://cloud.vespa.ai/en/getting-started-java">docs</a> for details.
|
||||
* Null if there is no SSL private key file e.g. for local Vespa server.
|
||||
* @param certPath local path to the SSL certificate file in PEM format. Read
|
||||
* <a href="https://cloud.vespa.ai/en/getting-started-java">docs</a> for details.
|
||||
* Null if there is no SSL certificate file e.g. for local Vespa server.
|
||||
* @param timeout for Vespa Java client in <code>java.time.Duration</code> format.
|
||||
* @param namespace required for document ID generation, find more details
|
||||
* <a href="https://docs.vespa.ai/en/documents.html#namespace">here</a>.
|
||||
* @param documentType document type, used for document ID generation, find more details
|
||||
* <a href="https://docs.vespa.ai/en/documents.html#namespace">here</a> and data querying
|
||||
* @param clusterName cluster name, used for deleting all documents, find more details
|
||||
* <a href="https://docs.vespa.ai/en/operations/batch-delete.html">here</a>
|
||||
* @param rankProfile rank profile from your .sd schema. Provided example schema configures cosine similarity match
|
||||
* @param targetHits sets the number of hits (10 is default) exposed to the real Vespa's first-phase ranking
|
||||
* function per content node, find more details
|
||||
* <a href="https://docs.vespa.ai/en/nearest-neighbor-search.html#querying-using-nearestneighbor-query-operator">here</a>.
|
||||
* @param avoidDups if true (default), then <code>VespaEmbeddingStore</code> will generate a hashed ID based on
|
||||
* provided text segment, which avoids duplicated entries in DB.
|
||||
* If false, then random ID will be generated.
|
||||
* @param logRequests If true, requests to the Vespa service are logged.
|
||||
* @param logResponses If true, responses from the Vespa service are logged.
|
||||
*/
|
||||
public VespaEmbeddingStore(
|
||||
String url,
|
||||
String keyPath,
|
||||
String certPath,
|
||||
Duration timeout,
|
||||
String namespace,
|
||||
String documentType,
|
||||
String clusterName,
|
||||
String rankProfile,
|
||||
Integer targetHits,
|
||||
Boolean avoidDups,
|
||||
Boolean logRequests,
|
||||
Boolean logResponses) {
|
||||
ensureNotNull(url, "url");
|
||||
|
||||
/**
|
||||
* Adds a new embedding with provided ID to the store.
|
||||
*
|
||||
* @param id "user-specified" part of document ID, find more details
|
||||
* <a href="https://docs.vespa.ai/en/documents.html#namespace">here</a>
|
||||
* @param embedding the embedding to add
|
||||
*/
|
||||
@Override
|
||||
public void add(String id, Embedding embedding) {
|
||||
add(id, embedding, null);
|
||||
}
|
||||
|
||||
@Override
|
||||
public String add(Embedding embedding, TextSegment textSegment) {
|
||||
return add(null, embedding, textSegment);
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<String> addAll(List<Embedding> embeddings) {
|
||||
return addAll(embeddings, null);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void addAll(List<String> ids, List<Embedding> embeddings, List<TextSegment> embedded) {
|
||||
if (embedded != null && embeddings.size() != embedded.size()) {
|
||||
throw new IllegalArgumentException("The list of embeddings and embedded must have the same size");
|
||||
this.url = url;
|
||||
this.keyPath = keyPath != null ? Paths.get(keyPath) : null;
|
||||
this.certPath = certPath != null ? Paths.get(certPath) : null;
|
||||
this.timeout = getOrDefault(timeout, DEFAULT_TIMEOUT);
|
||||
this.namespace = getOrDefault(namespace, DEFAULT_NAMESPACE);
|
||||
this.documentType = getOrDefault(documentType, DEFAULT_DOCUMENT_TYPE);
|
||||
this.clusterName = getOrDefault(clusterName, DEFAULT_CLUSTER_NAME);
|
||||
this.rankProfile = getOrDefault(rankProfile, DEFAULT_RANK_PROFILE);
|
||||
this.targetHits = getOrDefault(targetHits, DEFAULT_TARGET_HITS);
|
||||
this.avoidDups = getOrDefault(avoidDups, DEFAULT_AVOID_DUPS);
|
||||
this.logRequests = getOrDefault(logRequests, false);
|
||||
this.logResponses = getOrDefault(logResponses, false);
|
||||
}
|
||||
|
||||
try (JsonFeeder jsonFeeder = buildJsonFeeder()) {
|
||||
List<Record> records = new ArrayList<>();
|
||||
private static EmbeddingMatch<TextSegment> toEmbeddingMatch(Record in) {
|
||||
return new EmbeddingMatch<>(
|
||||
in.relevance(),
|
||||
DocumentId.of(in.fields().documentid()).userSpecific(),
|
||||
Embedding.from(in.fields().vector().values()),
|
||||
in.fields().textSegment() != null ? TextSegment.from(in.fields().textSegment()) : null);
|
||||
}
|
||||
|
||||
for (int i = 0; i < embeddings.size(); i++) {
|
||||
records.add(buildRecord(ids.get(i), embeddings.get(i), embedded != null ? embedded.get(i) : null));
|
||||
}
|
||||
public static Builder builder() {
|
||||
return new Builder();
|
||||
}
|
||||
|
||||
jsonFeeder.feedMany(
|
||||
Json.toInputStream(records, List.class),
|
||||
new JsonFeeder.ResultCallback() {
|
||||
@Override
|
||||
public void onNextResult(Result result, FeedException error) {
|
||||
if (error != null) {
|
||||
throw new RuntimeException(error.getMessage());
|
||||
}
|
||||
}
|
||||
@Override
|
||||
public void onError(FeedException error) {
|
||||
throw new RuntimeException(error.getMessage());
|
||||
}
|
||||
@Override
|
||||
public String add(Embedding embedding) {
|
||||
return add(null, embedding, null);
|
||||
}
|
||||
|
||||
/**
|
||||
* Adds a new embedding with provided ID to the store.
|
||||
*
|
||||
* @param id "user-specified" part of document ID, find more details
|
||||
* <a href="https://docs.vespa.ai/en/documents.html#namespace">here</a>
|
||||
* @param embedding the embedding to add
|
||||
*/
|
||||
@Override
|
||||
public void add(String id, Embedding embedding) {
|
||||
add(id, embedding, null);
|
||||
}
|
||||
|
||||
@Override
|
||||
public String add(Embedding embedding, TextSegment textSegment) {
|
||||
return add(null, embedding, textSegment);
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<String> addAll(List<Embedding> embeddings) {
|
||||
return addAll(embeddings, null);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void addAll(List<String> ids, List<Embedding> embeddings, List<TextSegment> embedded) {
|
||||
if (embedded != null && embeddings.size() != embedded.size()) {
|
||||
throw new IllegalArgumentException("The list of embeddings and embedded must have the same size");
|
||||
}
|
||||
);
|
||||
} catch (IOException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
|
||||
}
|
||||
try (JsonFeeder jsonFeeder = feeder()) {
|
||||
List<Record> records = new ArrayList<>();
|
||||
|
||||
/**
|
||||
* {@inheritDoc}
|
||||
* The score inside {@link EmbeddingMatch} is Vespa relevance according to provided rank profile.
|
||||
*/
|
||||
@Override
|
||||
@SneakyThrows
|
||||
public List<EmbeddingMatch<TextSegment>> findRelevant(Embedding referenceEmbedding, int maxResults, double minScore) {
|
||||
try {
|
||||
String searchQuery = Q
|
||||
.select(FIELD_NAME_DOCUMENT_ID, FIELD_NAME_TEXT_SEGMENT, FIELD_NAME_VECTOR)
|
||||
.from(documentType)
|
||||
.where(buildNearestNeighbor())
|
||||
.fix()
|
||||
.hits(maxResults)
|
||||
.ranking(rankProfile)
|
||||
.param("input.query(q)", Json.toJson(referenceEmbedding.vectorAsList()))
|
||||
.param("input.query(threshold)", String.valueOf(minScore))
|
||||
.build();
|
||||
|
||||
Response<QueryResponse> response = getQueryApi().search(searchQuery).execute();
|
||||
if (response.isSuccessful()) {
|
||||
QueryResponse parsedResponse = response.body();
|
||||
return parsedResponse
|
||||
.getRoot()
|
||||
.getChildren()
|
||||
.stream()
|
||||
.map(VespaEmbeddingStore::toEmbeddingMatch)
|
||||
.collect(Collectors.toList());
|
||||
} else {
|
||||
throw new RuntimeException("Request failed");
|
||||
}
|
||||
} catch (Exception e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
private String add(String id, Embedding embedding, TextSegment textSegment) {
|
||||
AtomicReference<String> resId = new AtomicReference<>();
|
||||
|
||||
try (JsonFeeder jsonFeeder = buildJsonFeeder()) {
|
||||
jsonFeeder
|
||||
.feedSingle(Json.toJson(buildRecord(id, embedding, textSegment)))
|
||||
.whenComplete(
|
||||
(
|
||||
(result, throwable) -> {
|
||||
if (throwable != null) {
|
||||
throw new RuntimeException(throwable);
|
||||
} else if (Result.Type.success.equals(result.type())) {
|
||||
resId.set(result.documentId().toString());
|
||||
}
|
||||
for (int i = 0; i < embeddings.size(); i++) {
|
||||
records.add(buildRecord(ids.get(i), embeddings.get(i), embedded != null ? embedded.get(i) : null));
|
||||
}
|
||||
)
|
||||
);
|
||||
} catch (Exception e) {
|
||||
throw new RuntimeException(e);
|
||||
|
||||
jsonFeeder.feedMany(
|
||||
new ByteArrayInputStream(
|
||||
OBJECT_MAPPER.writeValueAsString(records).getBytes()),
|
||||
new JsonFeeder.ResultCallback() {
|
||||
@Override
|
||||
public void onNextResult(Result result, FeedException error) {
|
||||
if (error != null) {
|
||||
throw new RuntimeException(error.getMessage());
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onError(FeedException error) {
|
||||
throw new RuntimeException(error.getMessage());
|
||||
}
|
||||
});
|
||||
} catch (IOException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
return resId.get();
|
||||
}
|
||||
/**
|
||||
* {@inheritDoc}
|
||||
* The score inside {@link EmbeddingMatch} is Vespa relevance according to provided rank profile.
|
||||
*/
|
||||
@Override
|
||||
public EmbeddingSearchResult<TextSegment> search(EmbeddingSearchRequest request) {
|
||||
try {
|
||||
String searchQuery = Q.select(FIELD_NAME_DOCUMENT_ID, FIELD_NAME_TEXT_SEGMENT, FIELD_NAME_VECTOR)
|
||||
.from(documentType)
|
||||
.where(buildNearestNeighbor())
|
||||
.fix()
|
||||
.hits(request.maxResults())
|
||||
.ranking(rankProfile)
|
||||
.param(
|
||||
"input.query(q)",
|
||||
OBJECT_MAPPER.writeValueAsString(
|
||||
request.queryEmbedding().vectorAsList()))
|
||||
.param("input.query(threshold)", String.valueOf(request.minScore()))
|
||||
.build();
|
||||
|
||||
private JsonFeeder buildJsonFeeder() {
|
||||
return JsonFeeder
|
||||
.builder(FeedClientBuilder.create(URI.create(url)).setCertificate(certPath, keyPath).build())
|
||||
.withTimeout(timeout)
|
||||
.build();
|
||||
}
|
||||
|
||||
private VespaQueryApi getQueryApi() {
|
||||
if (queryApi == null) {
|
||||
queryApi = createInstance(url, certPath, keyPath);
|
||||
Response<QueryResponse> response = api().search(searchQuery).execute();
|
||||
if (response.isSuccessful()) {
|
||||
QueryResponse parsedResponse = response.body();
|
||||
List<Record> children = parsedResponse.root().children();
|
||||
return new EmbeddingSearchResult<>(
|
||||
children == null || children.isEmpty()
|
||||
? new ArrayList<>()
|
||||
: children.stream()
|
||||
.map(VespaEmbeddingStore::toEmbeddingMatch)
|
||||
.toList());
|
||||
} else {
|
||||
throw toException(response);
|
||||
}
|
||||
} catch (Exception e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
return queryApi;
|
||||
}
|
||||
|
||||
private static EmbeddingMatch<TextSegment> toEmbeddingMatch(Record in) {
|
||||
return new EmbeddingMatch<>(
|
||||
in.getRelevance(),
|
||||
in.getFields().getDocumentId(),
|
||||
Embedding.from(in.getFields().getVector().getValues()),
|
||||
TextSegment.from(in.getFields().getTextSegment())
|
||||
);
|
||||
}
|
||||
@Override
|
||||
public void removeAll() {
|
||||
try {
|
||||
Response<DeleteResponse> response =
|
||||
api().deleteAll(namespace, documentType, clusterName).execute();
|
||||
if (!response.isSuccessful()) {
|
||||
throw toException(response);
|
||||
}
|
||||
} catch (Exception e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
private Record buildRecord(String id, Embedding embedding, TextSegment textSegment) {
|
||||
String recordId = id != null
|
||||
? id
|
||||
: avoidDups && textSegment != null ? generateUUIDFrom(textSegment.text()) : randomUUID();
|
||||
DocumentId documentId = DocumentId.of(namespace, documentType, recordId);
|
||||
String text = textSegment != null ? textSegment.text() : null;
|
||||
return new Record(documentId.toString(), text, embedding.vectorAsList());
|
||||
}
|
||||
private String add(String id, Embedding embedding, TextSegment textSegment) {
|
||||
AtomicReference<String> resId = new AtomicReference<>();
|
||||
|
||||
private Record buildRecord(Embedding embedding, TextSegment textSegment) {
|
||||
return buildRecord(null, embedding, textSegment);
|
||||
}
|
||||
try (JsonFeeder jsonFeeder = feeder()) {
|
||||
jsonFeeder
|
||||
.feedSingle(OBJECT_MAPPER.writeValueAsString(buildRecord(id, embedding, textSegment)))
|
||||
.whenComplete(((result, throwable) -> {
|
||||
if (throwable != null) {
|
||||
throw new RuntimeException(throwable);
|
||||
} else if (Result.Type.success.equals(result.type())) {
|
||||
resId.set(result.documentId().userSpecific());
|
||||
}
|
||||
}));
|
||||
} catch (Exception e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
|
||||
private NearestNeighbor buildNearestNeighbor()
|
||||
throws NoSuchMethodException, IllegalAccessException, InvocationTargetException {
|
||||
NearestNeighbor nb = Q.nearestNeighbor(FIELD_NAME_VECTOR, "q");
|
||||
return resId.get();
|
||||
}
|
||||
|
||||
// workaround to invoke ai.vespa.client.dsl.NearestNeighbor#annotate,
|
||||
// see https://github.com/vespa-engine/vespa/issues/28029
|
||||
// The bug is fixed in the meantime, but the baseline has been upgraded to Java 11, hence this workaround remains here
|
||||
Method method = NearestNeighbor.class.getDeclaredMethod("annotate", new Class<?>[] { Annotation.class });
|
||||
method.setAccessible(true);
|
||||
method.invoke(nb, A.a("targetHits", targetHits));
|
||||
return nb;
|
||||
}
|
||||
private JsonFeeder feeder() {
|
||||
FeedClientBuilder fcBuilder = FeedClientBuilder.create(URI.create(url));
|
||||
if (certPath != null && keyPath != null) {
|
||||
fcBuilder.setCertificate(certPath, keyPath);
|
||||
}
|
||||
|
||||
return JsonFeeder.builder(fcBuilder.build()).withTimeout(timeout).build();
|
||||
}
|
||||
|
||||
private VespaApi api() {
|
||||
if (api == null) {
|
||||
api = createInstance(url, certPath, keyPath, logRequests, logResponses);
|
||||
}
|
||||
return api;
|
||||
}
|
||||
|
||||
private Record buildRecord(String id, Embedding embedding, TextSegment textSegment) {
|
||||
String recordId = id != null
|
||||
? id
|
||||
: avoidDups && textSegment != null ? generateUUIDFrom(textSegment.text()) : randomUUID();
|
||||
DocumentId documentId = DocumentId.of(namespace, documentType, recordId);
|
||||
String text = textSegment != null ? textSegment.text() : null;
|
||||
return new Record(documentId.toString(), null, new Fields(null, text, new Vector(embedding.vectorAsList())));
|
||||
}
|
||||
|
||||
private NearestNeighbor buildNearestNeighbor() {
|
||||
NearestNeighbor nb = Q.nearestNeighbor(FIELD_NAME_VECTOR, "q");
|
||||
nb.annotate(A.a("targetHits", targetHits));
|
||||
return nb;
|
||||
}
|
||||
|
||||
public static class Builder {
|
||||
|
||||
private String url;
|
||||
private String keyPath;
|
||||
private String certPath;
|
||||
private Duration timeout;
|
||||
private String namespace;
|
||||
private String documentType;
|
||||
private String clusterName;
|
||||
private String rankProfile;
|
||||
private Integer targetHits;
|
||||
private Boolean avoidDups;
|
||||
private Boolean logRequests;
|
||||
private Boolean logResponses;
|
||||
|
||||
/**
|
||||
* @param url server url, local or cloud one. The latter you can find under Endpoint of your Vespa
|
||||
* application, e.g. https://alexey-heezer.langchain4j.mytenant346.aws-us-east-1c.dev.z.vespa-app.cloud/
|
||||
*/
|
||||
public Builder url(String url) {
|
||||
this.url = url;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @param keyPath local path to the SSL private key file in PEM format. Read
|
||||
* <a href="https://cloud.vespa.ai/en/getting-started-java">docs</a> for details.
|
||||
*/
|
||||
public Builder keyPath(String keyPath) {
|
||||
this.keyPath = keyPath;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @param certPath local path to the SSL certificate file in PEM format. Read
|
||||
* <a href="https://cloud.vespa.ai/en/getting-started-java">docs</a> for details.
|
||||
*/
|
||||
public Builder certPath(String certPath) {
|
||||
this.certPath = certPath;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @param timeout for Vespa Java client in <code>java.time.Duration</code> format.
|
||||
*/
|
||||
public Builder timeout(Duration timeout) {
|
||||
this.timeout = timeout;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @param namespace required for document ID generation, find more details
|
||||
* <a href="https://docs.vespa.ai/en/documents.html#namespace">here</a>.
|
||||
*/
|
||||
public Builder namespace(String namespace) {
|
||||
this.namespace = namespace;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @param documentType document type, used for document ID generation, find more details
|
||||
* <a href="https://docs.vespa.ai/en/documents.html#namespace">here</a> and data querying.
|
||||
*/
|
||||
public Builder documentType(String documentType) {
|
||||
this.documentType = documentType;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @param clusterName cluster name, used for deleting all documents, find more details
|
||||
* <a href="https://docs.vespa.ai/en/operations/batch-delete.html">here</a>
|
||||
*/
|
||||
public Builder clusterName(String clusterName) {
|
||||
this.clusterName = clusterName;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @param rankProfile rank profile from your .sd schema. Provided example schema configures cosine similarity match.
|
||||
*/
|
||||
public Builder rankProfile(String rankProfile) {
|
||||
this.rankProfile = rankProfile;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @param targetHits sets the number of hits (10 is default) exposed to the real Vespa's first-phase ranking
|
||||
* function per content node, find more details
|
||||
* <a href="https://docs.vespa.ai/en/nearest-neighbor-search.html#querying-using-nearestneighbor-query-operator">here</a>.
|
||||
*/
|
||||
public Builder targetHits(Integer targetHits) {
|
||||
this.targetHits = targetHits;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder avoidDups(Boolean avoidDups) {
|
||||
this.avoidDups = avoidDups;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder logRequests(Boolean logRequests) {
|
||||
this.logRequests = logRequests;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder logResponses(Boolean logResponses) {
|
||||
this.logResponses = logResponses;
|
||||
return this;
|
||||
}
|
||||
|
||||
public VespaEmbeddingStore build() {
|
||||
return new VespaEmbeddingStore(
|
||||
url,
|
||||
keyPath,
|
||||
certPath,
|
||||
timeout,
|
||||
namespace,
|
||||
documentType,
|
||||
rankProfile,
|
||||
clusterName,
|
||||
targetHits,
|
||||
avoidDups,
|
||||
logRequests,
|
||||
logResponses);
|
||||
}
|
||||
}
|
||||
|
||||
private static RuntimeException toException(Response<?> response) throws IOException {
|
||||
try (ResponseBody responseBody = response.errorBody()) {
|
||||
int code = response.code();
|
||||
if (responseBody != null) {
|
||||
String body = responseBody.string();
|
||||
String errorMessage = String.format("status code: %s; body: %s", code, body);
|
||||
return new RuntimeException(errorMessage);
|
||||
} else {
|
||||
return new RuntimeException(String.format("status code: %s;", code));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,11 +0,0 @@
|
|||
package dev.langchain4j.store.embedding.vespa;
|
||||
|
||||
import retrofit2.Call;
|
||||
import retrofit2.http.GET;
|
||||
import retrofit2.http.Path;
|
||||
|
||||
interface VespaQueryApi {
|
||||
|
||||
@GET("search/{query}")
|
||||
Call<QueryResponse> search(@Path(value = "query", encoded = true) String query);
|
||||
}
|
||||
|
|
@ -1,144 +0,0 @@
|
|||
// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
|
||||
package dev.langchain4j.store.embedding.vespa;
|
||||
|
||||
import com.google.gson.GsonBuilder;
|
||||
import dev.langchain4j.internal.Utils;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.security.GeneralSecurityException;
|
||||
import java.security.KeyFactory;
|
||||
import java.security.KeyStore;
|
||||
import java.security.PrivateKey;
|
||||
import java.security.cert.Certificate;
|
||||
import java.security.cert.X509Certificate;
|
||||
import java.security.spec.PKCS8EncodedKeySpec;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import javax.net.ssl.*;
|
||||
import okhttp3.HttpUrl;
|
||||
import okhttp3.OkHttpClient;
|
||||
import okhttp3.Request;
|
||||
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
|
||||
import org.bouncycastle.asn1.pkcs.PKCSObjectIdentifiers;
|
||||
import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
|
||||
import org.bouncycastle.asn1.x9.X9ObjectIdentifiers;
|
||||
import org.bouncycastle.cert.X509CertificateHolder;
|
||||
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
|
||||
import org.bouncycastle.jce.provider.BouncyCastleProvider;
|
||||
import org.bouncycastle.openssl.PEMKeyPair;
|
||||
import org.bouncycastle.openssl.PEMParser;
|
||||
import retrofit2.Retrofit;
|
||||
import retrofit2.converter.gson.GsonConverterFactory;
|
||||
|
||||
/**
|
||||
* This Workaround is needed because of <a href="https://github.com/vespa-engine/vespa/issues/28026">this request</a>.
|
||||
* It will be redundant as soon as vespa-client is implemented. This class is copied from <code>vespa-feed-client</code>.
|
||||
* BouncyCastle integration for creating a {@link SSLContext} instance from PEM encoded material
|
||||
*/
|
||||
class VespaQueryClient {
|
||||
|
||||
static final BouncyCastleProvider bcProvider = new BouncyCastleProvider();
|
||||
|
||||
public static VespaQueryApi createInstance(String baseUrl, Path certificate, Path privateKey) {
|
||||
try {
|
||||
KeyStore keystore = KeyStore.getInstance("PKCS12");
|
||||
keystore.load(null);
|
||||
keystore.setKeyEntry("cert", privateKey(privateKey), new char[0], certificates(certificate));
|
||||
// Protocol version must be equal to TlsContext.SSL_CONTEXT_VERSION or higher
|
||||
SSLContext sslContext = SSLContext.getInstance("TLSv1.3");
|
||||
sslContext.init(createKeyManagers(keystore), null, /*Default secure random algorithm*/null);
|
||||
|
||||
TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance(
|
||||
TrustManagerFactory.getDefaultAlgorithm()
|
||||
);
|
||||
trustManagerFactory.init(keystore);
|
||||
|
||||
OkHttpClient client = new OkHttpClient.Builder()
|
||||
.sslSocketFactory(sslContext.getSocketFactory(), (X509TrustManager) trustManagerFactory.getTrustManagers()[0])
|
||||
.addInterceptor(chain -> {
|
||||
// trick to format the query URL exactly how Vespa expects it (search/?query),
|
||||
// see https://docs.vespa.ai/en/reference/query-language-reference.html
|
||||
Request request = chain.request();
|
||||
HttpUrl url = request
|
||||
.url()
|
||||
.newBuilder()
|
||||
.removePathSegment(1)
|
||||
.addPathSegment("")
|
||||
.encodedQuery(request.url().encodedPathSegments().get(1))
|
||||
.build();
|
||||
request = request.newBuilder().url(url).build();
|
||||
return chain.proceed(request);
|
||||
})
|
||||
.build();
|
||||
|
||||
Retrofit retrofit = new Retrofit.Builder()
|
||||
.baseUrl(Utils.ensureTrailingForwardSlash(baseUrl))
|
||||
.client(client)
|
||||
.addConverterFactory(GsonConverterFactory.create(new GsonBuilder().create()))
|
||||
.build();
|
||||
|
||||
return retrofit.create(VespaQueryApi.class);
|
||||
} catch (Exception e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
private static KeyManager[] createKeyManagers(KeyStore keystore) throws GeneralSecurityException {
|
||||
KeyManagerFactory kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
|
||||
kmf.init(keystore, new char[0]);
|
||||
return kmf.getKeyManagers();
|
||||
}
|
||||
|
||||
private static Certificate[] certificates(Path file) throws IOException, GeneralSecurityException {
|
||||
try (PEMParser parser = new PEMParser(Files.newBufferedReader(file))) {
|
||||
List<X509Certificate> result = new ArrayList<>();
|
||||
Object pemObject;
|
||||
while ((pemObject = parser.readObject()) != null) {
|
||||
result.add(toX509Certificate(pemObject));
|
||||
}
|
||||
if (result.isEmpty()) throw new IOException("File contains no PEM encoded certificates: " + file);
|
||||
return result.toArray(new Certificate[0]);
|
||||
}
|
||||
}
|
||||
|
||||
private static PrivateKey privateKey(Path file) throws IOException, GeneralSecurityException {
|
||||
try (PEMParser parser = new PEMParser(Files.newBufferedReader(file))) {
|
||||
Object pemObject;
|
||||
while ((pemObject = parser.readObject()) != null) {
|
||||
if (pemObject instanceof PrivateKeyInfo) {
|
||||
PrivateKeyInfo keyInfo = (PrivateKeyInfo) pemObject;
|
||||
PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyInfo.getEncoded());
|
||||
return createKeyFactory(keyInfo).generatePrivate(keySpec);
|
||||
} else if (pemObject instanceof PEMKeyPair) {
|
||||
PEMKeyPair pemKeypair = (PEMKeyPair) pemObject;
|
||||
PrivateKeyInfo keyInfo = pemKeypair.getPrivateKeyInfo();
|
||||
return createKeyFactory(keyInfo).generatePrivate(new PKCS8EncodedKeySpec(keyInfo.getEncoded()));
|
||||
}
|
||||
}
|
||||
throw new IOException("Could not find private key in PEM file");
|
||||
}
|
||||
}
|
||||
|
||||
private static X509Certificate toX509Certificate(Object pemObject) throws IOException, GeneralSecurityException {
|
||||
if (pemObject instanceof X509Certificate) return (X509Certificate) pemObject;
|
||||
if (pemObject instanceof X509CertificateHolder) {
|
||||
return new JcaX509CertificateConverter()
|
||||
.setProvider(bcProvider)
|
||||
.getCertificate((X509CertificateHolder) pemObject);
|
||||
}
|
||||
throw new IOException("Invalid type of PEM object: " + pemObject);
|
||||
}
|
||||
|
||||
private static KeyFactory createKeyFactory(PrivateKeyInfo info) throws IOException, GeneralSecurityException {
|
||||
ASN1ObjectIdentifier algorithm = info.getPrivateKeyAlgorithm().getAlgorithm();
|
||||
if (X9ObjectIdentifiers.id_ecPublicKey.equals(algorithm)) {
|
||||
return KeyFactory.getInstance("EC", bcProvider);
|
||||
} else if (PKCSObjectIdentifiers.rsaEncryption.equals(algorithm)) {
|
||||
return KeyFactory.getInstance("RSA", bcProvider);
|
||||
} else {
|
||||
throw new IOException("Unknown key algorithm: " + algorithm);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,50 @@
|
|||
package dev.langchain4j.store.embedding.vespa;
|
||||
|
||||
import static dev.langchain4j.store.embedding.vespa.VespaResponseLoggingInterceptor.getHeaders;
|
||||
|
||||
import java.io.IOException;
|
||||
import okhttp3.Interceptor;
|
||||
import okhttp3.Request;
|
||||
import okhttp3.Response;
|
||||
import okio.Buffer;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
|
||||
class VespaRequestLoggingInterceptor implements Interceptor {
|
||||
|
||||
private static final Logger log = LoggerFactory.getLogger(VespaRequestLoggingInterceptor.class);
|
||||
|
||||
@Override
|
||||
public Response intercept(Interceptor.Chain chain) throws IOException {
|
||||
Request request = chain.request();
|
||||
this.log(request);
|
||||
return chain.proceed(request);
|
||||
}
|
||||
|
||||
private void log(Request request) {
|
||||
try {
|
||||
log.debug(
|
||||
"Request:\n- method: {}\n- url: {}\n- headers: {}\n- body: {}",
|
||||
request.method(),
|
||||
request.url(),
|
||||
getHeaders(request.headers()),
|
||||
getBody(request));
|
||||
} catch (Exception e) {
|
||||
log.warn("Error while logging request: {}", e.getMessage());
|
||||
}
|
||||
}
|
||||
|
||||
private static String getBody(Request request) {
|
||||
try {
|
||||
Buffer buffer = new Buffer();
|
||||
if (request.body() == null) {
|
||||
return "";
|
||||
}
|
||||
request.body().writeTo(buffer);
|
||||
return buffer.readUtf8();
|
||||
} catch (Exception e) {
|
||||
log.warn("Exception while getting body", e);
|
||||
return "Exception while getting body: " + e.getMessage();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,50 @@
|
|||
package dev.langchain4j.store.embedding.vespa;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.util.stream.Collectors;
|
||||
import java.util.stream.StreamSupport;
|
||||
import okhttp3.Headers;
|
||||
import okhttp3.Interceptor;
|
||||
import okhttp3.Request;
|
||||
import okhttp3.Response;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
|
||||
class VespaResponseLoggingInterceptor implements Interceptor {
|
||||
|
||||
private static final Logger log = LoggerFactory.getLogger(VespaResponseLoggingInterceptor.class);
|
||||
|
||||
@Override
|
||||
public Response intercept(Chain chain) throws IOException {
|
||||
Request request = chain.request();
|
||||
Response response = chain.proceed(request);
|
||||
this.log(response);
|
||||
return response;
|
||||
}
|
||||
|
||||
private void log(Response response) {
|
||||
try {
|
||||
log.debug(
|
||||
"Response:\n- status code: {}\n- headers: {}\n- body: {}",
|
||||
response.code(),
|
||||
getHeaders(response.headers()),
|
||||
getBody(response));
|
||||
} catch (Exception e) {
|
||||
log.warn("Error while logging response: {}", e.getMessage());
|
||||
}
|
||||
}
|
||||
|
||||
private static String getBody(Response response) throws IOException {
|
||||
return response.peekBody(Long.MAX_VALUE).string();
|
||||
}
|
||||
|
||||
static String getHeaders(Headers headers) {
|
||||
return StreamSupport.stream(headers.spliterator(), false)
|
||||
.map(header -> {
|
||||
String headerKey = header.component1();
|
||||
String headerValue = header.component2();
|
||||
return String.format("[%s: %s]", headerKey, headerValue);
|
||||
})
|
||||
.collect(Collectors.joining(", "));
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,35 @@
|
|||
package dev.langchain4j.store.embedding.vespa;
|
||||
|
||||
import dev.langchain4j.data.segment.TextSegment;
|
||||
import dev.langchain4j.model.embedding.EmbeddingModel;
|
||||
import dev.langchain4j.model.embedding.onnx.allminilml6v2q.AllMiniLmL6V2QuantizedEmbeddingModel;
|
||||
import dev.langchain4j.store.embedding.EmbeddingStore;
|
||||
import dev.langchain4j.store.embedding.EmbeddingStoreIT;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariables;
|
||||
|
||||
@EnabledIfEnvironmentVariables({
|
||||
@EnabledIfEnvironmentVariable(named = "VESPA_URL", matches = ".+"),
|
||||
@EnabledIfEnvironmentVariable(named = "VESPA_KEY_PATH", matches = ".+"),
|
||||
@EnabledIfEnvironmentVariable(named = "VESPA_CERT_PATH", matches = ".+")
|
||||
})
|
||||
public class VespaEmbeddingStoreCloudIT extends EmbeddingStoreIT {
|
||||
|
||||
EmbeddingStore<TextSegment> embeddingStore = VespaEmbeddingStore.builder()
|
||||
.url(System.getenv("VESPA_URL"))
|
||||
.keyPath(System.getenv("VESPA_KEY_PATH"))
|
||||
.certPath(System.getenv("VESPA_CERT_PATH"))
|
||||
.build();
|
||||
|
||||
EmbeddingModel embeddingModel = new AllMiniLmL6V2QuantizedEmbeddingModel();
|
||||
|
||||
@Override
|
||||
protected EmbeddingStore<TextSegment> embeddingStore() {
|
||||
return embeddingStore;
|
||||
}
|
||||
|
||||
@Override
|
||||
protected EmbeddingModel embeddingModel() {
|
||||
return embeddingModel;
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,150 @@
|
|||
package dev.langchain4j.store.embedding.vespa;
|
||||
|
||||
import static java.util.Arrays.asList;
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.awaitility.Awaitility.await;
|
||||
|
||||
import dev.langchain4j.data.embedding.Embedding;
|
||||
import dev.langchain4j.data.segment.TextSegment;
|
||||
import dev.langchain4j.model.embedding.EmbeddingModel;
|
||||
import dev.langchain4j.model.embedding.onnx.allminilml6v2q.AllMiniLmL6V2QuantizedEmbeddingModel;
|
||||
import dev.langchain4j.store.embedding.EmbeddingMatch;
|
||||
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
|
||||
import dev.langchain4j.store.embedding.EmbeddingStore;
|
||||
import dev.langchain4j.store.embedding.EmbeddingStoreWithoutMetadataIT;
|
||||
import java.io.DataOutputStream;
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.net.HttpURLConnection;
|
||||
import java.net.URI;
|
||||
import java.net.URISyntaxException;
|
||||
import java.time.Duration;
|
||||
import java.util.List;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
import org.testcontainers.junit.jupiter.Container;
|
||||
import org.testcontainers.junit.jupiter.Testcontainers;
|
||||
import org.testcontainers.utility.DockerImageName;
|
||||
|
||||
@Testcontainers
|
||||
class VespaEmbeddingStoreIT extends EmbeddingStoreWithoutMetadataIT {
|
||||
|
||||
@Container
|
||||
static GenericContainer<?> vespa = new GenericContainer<>(DockerImageName.parse("vespaengine/vespa:8.458.13"))
|
||||
.waitingFor(Wait.forListeningPorts(19071))
|
||||
.withExposedPorts(8080, 19071);
|
||||
|
||||
private final EmbeddingModel embeddingModel = new AllMiniLmL6V2QuantizedEmbeddingModel();
|
||||
|
||||
private final VespaEmbeddingStore embeddingStore = VespaEmbeddingStore.builder()
|
||||
.url(String.format("http://%s:%d", vespa.getHost(), vespa.getMappedPort(8080)))
|
||||
.build();
|
||||
|
||||
@Test
|
||||
void should_find_relevant_matches() {
|
||||
// given
|
||||
TextSegment segment1 = TextSegment.from("I like football.");
|
||||
Embedding embedding1 = embeddingModel.embed(segment1).content();
|
||||
|
||||
TextSegment segment2 = TextSegment.from("I've never been to New York.");
|
||||
Embedding embedding2 = embeddingModel.embed(segment2).content();
|
||||
|
||||
TextSegment segment3 =
|
||||
TextSegment.from("But actually we tried our new swimming pool yesterday and it was awesome!");
|
||||
Embedding embedding3 = embeddingModel.embed(segment3).content();
|
||||
|
||||
TextSegment segment4 = TextSegment.from("John Lennon was a very cool person.");
|
||||
Embedding embedding4 = embeddingModel.embed(segment4).content();
|
||||
|
||||
embeddingStore.addAll(
|
||||
asList(embedding1, embedding2, embedding3, embedding4), asList(segment1, segment2, segment3, segment4));
|
||||
|
||||
// when
|
||||
Embedding sportEmbedding =
|
||||
embeddingModel.embed("What is your favorite sport?").content();
|
||||
List<EmbeddingMatch<TextSegment>> sportMatches = embeddingStore
|
||||
.search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(sportEmbedding)
|
||||
.maxResults(2)
|
||||
.build())
|
||||
.matches();
|
||||
Embedding musicEmbedding =
|
||||
embeddingModel.embed("And what about musicians?").content();
|
||||
List<EmbeddingMatch<TextSegment>> musicMatches = embeddingStore
|
||||
.search(EmbeddingSearchRequest.builder()
|
||||
.queryEmbedding(musicEmbedding)
|
||||
.maxResults(5)
|
||||
.minScore(0.6)
|
||||
.build())
|
||||
.matches();
|
||||
|
||||
// then
|
||||
assertThat(sportMatches).hasSize(2);
|
||||
assertThat(sportMatches.get(0).score()).isCloseTo(0.808, percentage());
|
||||
assertThat(sportMatches.get(0).embedded().text()).contains("football");
|
||||
assertThat(sportMatches.get(1).score()).isCloseTo(0.606, percentage());
|
||||
assertThat(sportMatches.get(1).embedded().text()).contains("swimming pool");
|
||||
|
||||
assertThat(musicMatches).hasSize(1);
|
||||
assertThat(musicMatches.get(0).score()).isCloseTo(0.671, percentage());
|
||||
assertThat(musicMatches.get(0).embedded().text()).contains("John Lennon");
|
||||
}
|
||||
|
||||
@Override
|
||||
protected EmbeddingStore<TextSegment> embeddingStore() {
|
||||
return embeddingStore;
|
||||
}
|
||||
|
||||
@Override
|
||||
protected EmbeddingModel embeddingModel() {
|
||||
return embeddingModel;
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void clearStore() {
|
||||
embeddingStore().removeAll();
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void ensureStoreIsEmpty() {
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(super::ensureStoreIsEmpty);
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void ensureStoreIsReady() {
|
||||
try {
|
||||
getAllEmbeddings();
|
||||
} catch (Exception e) {
|
||||
deployVespaApp();
|
||||
}
|
||||
}
|
||||
|
||||
private void deployVespaApp() {
|
||||
try {
|
||||
URI uri = new URI(String.format(
|
||||
"http://%s:%d/application/v2/tenant/default/prepareandactivate",
|
||||
vespa.getHost(), vespa.getMappedPort(19071)));
|
||||
HttpURLConnection conn = (HttpURLConnection) uri.toURL().openConnection();
|
||||
conn.setRequestMethod("POST");
|
||||
conn.setDoOutput(true);
|
||||
conn.setRequestProperty("Content-Type", "application/zip");
|
||||
|
||||
try (InputStream in = VespaEmbeddingStoreIT.class.getResourceAsStream("/vespa_app.zip");
|
||||
DataOutputStream out = new DataOutputStream(conn.getOutputStream())) {
|
||||
byte[] buffer = new byte[8192];
|
||||
int read;
|
||||
while ((read = in.read(buffer)) != -1) {
|
||||
out.write(buffer, 0, read);
|
||||
}
|
||||
}
|
||||
|
||||
assertThat(conn.getResponseCode()).isEqualTo(200);
|
||||
|
||||
// wait for Vespa application is deployed properly
|
||||
await().atMost(Duration.ofSeconds(60)).ignoreExceptions().untilAsserted(super::ensureStoreIsEmpty);
|
||||
} catch (URISyntaxException | IOException e) {
|
||||
e.printStackTrace();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1 @@
|
|||
writer.level = info
|
||||
Binary file not shown.
Loading…
Reference in New Issue