Improve the Document API

- Add ContentFormatter and DefaultContentFormatter that can filter the metadata
   and format the Document metadata and text according to predefined templates.
 - Add content formatter tests
 - Allow the TextSplitter to copy the document  content-formatter to the children.
    When the splitter breaks the parent Document into multiple chunks (e.g.
    into a list of children Documents) copy the source content formatter to
    the chunks by default. Use the copyContentFormatter flag to enable/disable copping.
 - Add TextSplitter IT tests
 - Add MetadataExtractors as DocumentTransformers.
 - Bump spring-ai project version to 0.7.0-SNAPSHOT
 - Configurable metadata-mode for EmbeddingClients
   - Make the metadata mode configurable for the EmbeddingClient implementations.
   - Use the EMBED mode by default.

 Resolves #44
This commit is contained in:
Christian Tzolov
2023-09-20 09:56:18 +02:00
committed by Mark Pollack
parent 5c23b9d3d8
commit 76e44d68e4
39 changed files with 1290 additions and 154 deletions

View File

@@ -37,7 +37,7 @@ And the Spring Boot Starter depending on if you are using Azure Open AI or Open
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-azure-openai-spring-boot-starter</artifactId>
<version>0.2.0-SNAPSHOT</version>
<version>0.7.0-SNAPSHOT</version>
</dependency>
```
@@ -47,7 +47,7 @@ And the Spring Boot Starter depending on if you are using Azure Open AI or Open
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
<version>0.2.0-SNAPSHOT</version>
<version>0.7.0-SNAPSHOT</version>
</dependency>
```
@@ -59,7 +59,7 @@ And the Spring Boot Starter depending on if you are using Azure Open AI or Open
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
<version>0.2.0-SNAPSHOT</version>
<version>0.7.0-SNAPSHOT</version>
</dependency>
```

View File

@@ -3,7 +3,7 @@
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.2.0-SNAPSHOT</version>
<version>0.7.0-SNAPSHOT</version>
<packaging>pom</packaging>
<url>https://github.com/spring-projects-experimental/spring-ai</url>

View File

@@ -4,7 +4,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.2.0-SNAPSHOT</version>
<version>0.7.0-SNAPSHOT</version>
</parent>
<artifactId>spring-ai-azure-openai</artifactId>
<packaging>jar</packaging>

View File

@@ -15,6 +15,7 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.embedding.Embedding;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.embedding.EmbeddingResponse;
@@ -31,15 +32,23 @@ public class AzureOpenAiEmbeddingClient implements EmbeddingClient {
private final AtomicInteger embeddingDimensions = new AtomicInteger(-1);
private final MetadataMode metadataMode;
public AzureOpenAiEmbeddingClient(OpenAIClient azureOpenAiClient) {
this(azureOpenAiClient, "text-embedding-ada-002");
}
public AzureOpenAiEmbeddingClient(OpenAIClient azureOpenAiClient, String model) {
this(azureOpenAiClient, model, MetadataMode.EMBED);
}
public AzureOpenAiEmbeddingClient(OpenAIClient azureOpenAiClient, String model, MetadataMode metadataMode) {
Assert.notNull(azureOpenAiClient, "com.azure.ai.openai.OpenAIClient must not be null");
Assert.notNull(model, "Model must not be null");
Assert.notNull(metadataMode, "Metadata mode must not be null");
this.azureOpenAiClient = azureOpenAiClient;
this.model = model;
this.metadataMode = metadataMode;
}
@Override
@@ -54,7 +63,7 @@ public class AzureOpenAiEmbeddingClient implements EmbeddingClient {
public List<Double> embed(Document document) {
logger.debug("Retrieving embeddings");
Embeddings embeddings = this.azureOpenAiClient.getEmbeddings(this.model,
new EmbeddingsOptions(List.of(document.getContent())));
new EmbeddingsOptions(List.of(document.getFormattedContent(this.metadataMode))));
logger.debug("Embeddings retrieved");
return extractEmbeddingsList(embeddings);
}

View File

@@ -5,7 +5,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.2.0-SNAPSHOT</version>
<version>0.7.0-SNAPSHOT</version>
</parent>
<artifactId>spring-ai-core</artifactId>
<packaging>jar</packaging>

View File

@@ -55,7 +55,7 @@ public class AiResponse {
}
/**
* Arbitrary LLM-provider specific output
* Arbitrary model provider specific output
*/
public Map<String, Object> getProviderOutput() {
return Collections.unmodifiableMap(providerOutput);

View File

@@ -0,0 +1,28 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.document;
/**
* Converts the Document text and metadata into a AI, prompt-friendly text representation.
*
* @author Christian Tzolov
*/
public interface ContentFormatter {
String format(Document document, MetadataMode mode);
}

View File

@@ -0,0 +1,266 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.document;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
import org.springframework.util.Assert;
/**
* @author Christian Tzolov
*/
public class DefaultContentFormatter implements ContentFormatter {
private static final String TEMPLATE_CONTENT_PLACEHOLDER = "{content}";
private static final String TEMPLATE_METADATA_STRING_PLACEHOLDER = "{metadata_string}";
private static final String TEMPLATE_VALUE_PLACEHOLDER = "{value}";
private static final String TEMPLATE_KEY_PLACEHOLDER = "{key}";
private static final String DEFAULT_METADATA_TEMPLATE = String.format("%s: %s", TEMPLATE_KEY_PLACEHOLDER,
TEMPLATE_VALUE_PLACEHOLDER);
private static final String DEFAULT_METADATA_SEPARATOR = "\n";
private static final String DEFAULT_TEXT_TEMPLATE = String.format("%s\n\n%s", TEMPLATE_METADATA_STRING_PLACEHOLDER,
TEMPLATE_CONTENT_PLACEHOLDER);
/**
* Template for how metadata is formatted, with {key} and {value} placeholders.
*/
private final String metadataTemplate;
/**
* Separator between metadata fields when converting to string.
*/
private final String metadataSeparator;
/**
* Template for how Document text is formatted, with {content} and {metadata_string}
* placeholders.
*/
private final String textTemplate;
/**
* Metadata keys that are excluded from text for the inference.
*/
private final List<String> excludedInferenceMetadataKeys;
/**
* Metadata keys that are excluded from text for the embed model.
*/
private final List<String> excludedEmbedMetadataKeys;
/**
* Start building a new configuration.
* @return The entry point for creating a new configuration.
*/
public static Builder builder() {
return new Builder();
}
/**
* {@return the default config}
*/
public static DefaultContentFormatter defaultConfig() {
return builder().build();
}
private DefaultContentFormatter(Builder builder) {
this.metadataTemplate = builder.metadataTemplate;
this.metadataSeparator = builder.metadataSeparator;
this.textTemplate = builder.textTemplate;
this.excludedInferenceMetadataKeys = builder.excludedInferenceMetadataKeys;
this.excludedEmbedMetadataKeys = builder.excludedEmbedMetadataKeys;
}
public static class Builder {
private String metadataTemplate = DEFAULT_METADATA_TEMPLATE;
private String metadataSeparator = DEFAULT_METADATA_SEPARATOR;
private String textTemplate = DEFAULT_TEXT_TEMPLATE;
private List<String> excludedInferenceMetadataKeys = new ArrayList<>();
private List<String> excludedEmbedMetadataKeys = new ArrayList<>();
private Builder() {
}
public Builder from(DefaultContentFormatter fromFormatter) {
this.withExcludedEmbedMetadataKeys(fromFormatter.getExcludedEmbedMetadataKeys())
.withExcludedInferenceMetadataKeys(fromFormatter.getExcludedInferenceMetadataKeys())
.withMetadataSeparator(fromFormatter.getMetadataSeparator())
.withMetadataTemplate(fromFormatter.getMetadataTemplate())
.withTextTemplate(fromFormatter.getTextTemplate());
return this;
}
/**
* Configures the Document metadata template.
* @param metadataTemplate Metadata template to use.
* @return this builder
*/
public Builder withMetadataTemplate(String metadataTemplate) {
Assert.hasText(metadataTemplate, "Metadata Template must not be empty");
this.metadataTemplate = metadataTemplate;
return this;
}
/**
* Configures the Document metadata separator.
* @param metadataSeparator Metadata separator to use.
* @return this builder
*/
public Builder withMetadataSeparator(String metadataSeparator) {
Assert.notNull(metadataSeparator, "Metadata separator must not be empty");
this.metadataSeparator = metadataSeparator;
return this;
}
/**
* Configures the Document text template.
* @param textTemplate Document's content template.
* @return this builder
*/
public Builder withTextTemplate(String textTemplate) {
Assert.hasText(textTemplate, "Document's text template must not be empty");
this.textTemplate = textTemplate;
return this;
}
/**
* Configures the excluded Inference metadata keys to filter out from the model.
* @param excludedInferenceMetadataKeys Excluded inference metadata keys to use.
* @return this builder
*/
public Builder withExcludedInferenceMetadataKeys(List<String> excludedInferenceMetadataKeys) {
Assert.notNull(excludedInferenceMetadataKeys, "Excluded inference metadata keys must not be null");
this.excludedInferenceMetadataKeys = excludedInferenceMetadataKeys;
return this;
}
public Builder withExcludedInferenceMetadataKeys(String... keys) {
Assert.notNull(keys, "Excluded inference metadata keys must not be null");
this.excludedInferenceMetadataKeys.addAll(Arrays.asList(keys));
return this;
}
/**
* Configures the excluded Embed metadata keys to filter out from the model.
* @param excludedEmbedMetadataKeys Excluded Embed metadata keys to use.
* @return this builder
*/
public Builder withExcludedEmbedMetadataKeys(List<String> excludedEmbedMetadataKeys) {
Assert.notNull(excludedEmbedMetadataKeys, "Excluded Embed metadata keys must not be null");
this.excludedEmbedMetadataKeys = excludedEmbedMetadataKeys;
return this;
}
public Builder withExcludedEmbedMetadataKeys(String... keys) {
Assert.notNull(keys, "Excluded Embed metadata keys must not be null");
this.excludedEmbedMetadataKeys.addAll(Arrays.asList(keys));
return this;
}
/**
* {@return the immutable configuration}
*/
public DefaultContentFormatter build() {
return new DefaultContentFormatter(this);
}
}
@Override
public String format(Document document, MetadataMode metadataMode) {
var metadata = metadataFilter(document.getMetadata(), metadataMode);
var metadataText = metadata.entrySet()
.stream()
.map(metadataEntry -> this.metadataTemplate.replace(TEMPLATE_KEY_PLACEHOLDER, metadataEntry.getKey())
.replace(TEMPLATE_VALUE_PLACEHOLDER, metadataEntry.getValue().toString()))
.collect(Collectors.joining(this.metadataSeparator));
return this.textTemplate.replace(TEMPLATE_METADATA_STRING_PLACEHOLDER, metadataText)
.replace(TEMPLATE_CONTENT_PLACEHOLDER, document.getContent());
}
/**
* Filters the metadata by the configured MetadataMode.
* @param metadata Document metadata.
* @return Returns the filtered by configured mode metadata.
*/
protected Map<String, Object> metadataFilter(Map<String, Object> metadata, MetadataMode metadataMode) {
if (metadataMode == MetadataMode.ALL) {
return new HashMap<String, Object>(metadata);
}
if (metadataMode == MetadataMode.NONE) {
return new HashMap<String, Object>(Collections.emptyMap());
}
Set<String> usableMetadataKeys = new HashSet<>(metadata.keySet());
if (metadataMode == MetadataMode.INFERENCE) {
usableMetadataKeys.removeAll(this.excludedInferenceMetadataKeys);
}
else if (metadataMode == MetadataMode.EMBED) {
usableMetadataKeys.removeAll(this.excludedEmbedMetadataKeys);
}
return new HashMap<String, Object>(metadata.entrySet()
.stream()
.filter(e -> usableMetadataKeys.contains(e.getKey()))
.collect(Collectors.toMap(e -> e.getKey(), e -> e.getValue())));
}
public String getMetadataTemplate() {
return this.metadataTemplate;
}
public String getMetadataSeparator() {
return this.metadataSeparator;
}
public String getTextTemplate() {
return this.textTemplate;
}
public List<String> getExcludedInferenceMetadataKeys() {
return Collections.unmodifiableList(this.excludedInferenceMetadataKeys);
}
public List<String> getExcludedEmbedMetadataKeys() {
return Collections.unmodifiableList(this.excludedEmbedMetadataKeys);
}
}

View File

@@ -1,44 +1,83 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.document;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonProperty;
import org.springframework.util.StringUtils;
import java.util.*;
import org.springframework.util.Assert;
@JsonIgnoreProperties({ "contentFormatter" })
public class Document {
public final static ContentFormatter DEFAULT_CONTENT_FORMATTER = DefaultContentFormatter.defaultConfig();
/**
* Unique ID
*/
private final String id;
@JsonProperty(index = 100)
private List<Double> embedding = new ArrayList<>();
/**
* Metadata for the document. It should not be nested and values should be restricted
* to string, int, float, boolean for simple use with Vector Dbs.
*/
private Map<String, Object> metadata;
// Type; introduce when support images, now only text.
/**
* Document content.
*/
private String content;
private String text;
/**
* Embedding of the document. Note: ephemeral field.
*/
@JsonProperty(index = 100)
private List<Double> embedding = new ArrayList<>();
/**
* Mutable, ephemeral, content to text formatter. Defaults to Document text.
*/
@JsonIgnore
private ContentFormatter contentFormatter = DEFAULT_CONTENT_FORMATTER;
@JsonCreator(mode = JsonCreator.Mode.PROPERTIES)
public Document(@JsonProperty("text") String text) {
this(text, new HashMap<>());
public Document(@JsonProperty("content") String content) {
this(content, new HashMap<>());
}
public Document(String text, Map<String, Object> metadata) {
this(UUID.randomUUID().toString(), text, metadata);
public Document(String content, Map<String, Object> metadata) {
this(UUID.randomUUID().toString(), content, metadata);
}
public Document(String id, String text, Map<String, Object> metadata) {
public Document(String id, String content, Map<String, Object> metadata) {
Assert.hasText(id, "id must not be null");
Assert.hasText(content, "content must not be null");
Assert.notNull(metadata, "metadata must not be null");
this.id = id;
this.text = text;
this.content = content;
this.metadata = metadata;
}
@@ -46,109 +85,98 @@ public class Document {
return id;
}
public String getText() {
return this.text;
public String getContent() {
return this.content;
}
public Map<String, Object> getMetadata() {
return metadata;
@JsonIgnore
public String getFormattedContent() {
return this.getFormattedContent(MetadataMode.ALL);
}
public List<Double> getEmbedding() {
return embedding;
public String getFormattedContent(MetadataMode metadataMode) {
Assert.notNull(metadataMode, "Metadata mode must not be null");
return this.contentFormatter.format(this, metadataMode);
}
/**
* Helper content extractor that uses and external {@link ContentFormatter}.
*/
public String getFormattedContent(ContentFormatter formatter, MetadataMode metadataMode) {
Assert.notNull(formatter, "formatter must not be null");
Assert.notNull(metadataMode, "Metadata mode must not be null");
return formatter.format(this, metadataMode);
}
public void setEmbedding(List<Double> embedding) {
Assert.notNull(embedding, "embedding must not be null");
this.embedding = embedding;
}
/**
* Replace the document's {@link ContentFormatter}.
* @param contentFormatter new formatter to use.
*/
public void setContentFormatter(ContentFormatter contentFormatter) {
this.contentFormatter = contentFormatter;
}
public Map<String, Object> getMetadata() {
return this.metadata;
}
public List<Double> getEmbedding() {
return this.embedding;
}
public ContentFormatter getContentFormatter() {
return contentFormatter;
}
@Override
public int hashCode() {
final int prime = 31;
int result = 1;
result = prime * result + ((id == null) ? 0 : id.hashCode());
result = prime * result + ((metadata == null) ? 0 : metadata.hashCode());
result = prime * result + ((content == null) ? 0 : content.hashCode());
return result;
}
@Override
public boolean equals(Object obj) {
if (this == obj)
return true;
if (obj == null)
return false;
if (getClass() != obj.getClass())
return false;
Document other = (Document) obj;
if (id == null) {
if (other.id != null)
return false;
}
else if (!id.equals(other.id))
return false;
if (metadata == null) {
if (other.metadata != null)
return false;
}
else if (!metadata.equals(other.metadata))
return false;
if (content == null) {
if (other.content != null)
return false;
}
else if (!content.equals(other.content))
return false;
return true;
}
@Override
public String toString() {
return "Document{" + "id='" + id + '\'' + ", metadata=" + metadata + ", text='" + text + '\'' + '}';
}
private static String DEFAULT_TEXT_TEMPLATE = "{metadata_string}\n\n{text}";
private static String DEFAULT_METADATA_TEMPLATE = "{key}: {value}";
private String textTemplate = DEFAULT_TEXT_TEMPLATE;
private String metadataTemplate = DEFAULT_METADATA_TEMPLATE;
private String metadataSeparator = "\n";
private MetadataMode metadataMode = MetadataMode.NONE;
private List<String> excludedMetadataKeysForLlm;
@JsonIgnore
public String getContent() {
return getContent(MetadataMode.ALL);
}
public String getContent(MetadataMode metadataMode) {
if (metadataMode == MetadataMode.NONE) {
return this.text;
}
String metadataString = getMetadataString(metadataMode);
if (!StringUtils.hasText(metadataString)) {
return this.text;
}
return getTextTemplate().replace("{metadata_string}", metadataString).replace("{text}", text);
}
@JsonIgnore
public String getMetadataString() {
return getMetadataString(metadataMode);
}
public String getMetadataString(MetadataMode metadataMode) {
if (metadataMode == MetadataMode.NONE) {
return "";
}
Set<String> usableMetadataKeys = new HashSet<>(metadata.keySet());
if (metadataMode == MetadataMode.LLM) {
usableMetadataKeys.removeAll(this.excludedMetadataKeysForLlm);
}
else if (metadataMode == MetadataMode.EMBED) {
usableMetadataKeys.removeAll(this.excludedMetadataKeysForLlm);
}
List<String> metadataStringList = new ArrayList<>();
for (Map.Entry<String, Object> entry : metadata.entrySet()) {
String key = entry.getKey();
Object value = entry.getValue();
if (usableMetadataKeys.contains(key)) {
metadataStringList
.add(getMetadataTemplate().replace("{key}", key).replace("{value}", value.toString()));
}
}
return String.join(getMetadataSeparator(), metadataStringList);
}
private String getTextTemplate() {
return textTemplate;
}
private String getMetadataTemplate() {
return metadataTemplate;
}
private String getMetadataSeparator() {
return metadataSeparator;
}
public void setTextTemplate(String textTemplate) {
this.textTemplate = textTemplate;
}
public void setMetadataTemplate(String metadataTemplate) {
this.metadataTemplate = metadataTemplate;
}
public void setMetadataSeparator(String metadataSeparator) {
this.metadataSeparator = metadataSeparator;
return "Document{" + "id='" + id + '\'' + ", metadata=" + metadata + ", content='" + new String(content) + '\''
+ '}';
}
}

View File

@@ -1,3 +1,19 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.document;
import java.util.List;

View File

@@ -1,7 +1,23 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.document;
public enum MetadataMode {
ALL, EMBED, LLM, NONE;
ALL, EMBED, INFERENCE, NONE;
}
}

View File

@@ -0,0 +1,49 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.loader.extractor;
import java.util.List;
import java.util.Map;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentTransformer;
/**
* @author Christian Tzolov
*/
public abstract class AbstractMetadataFeatureExtractor implements DocumentTransformer {
@Override
public List<Document> apply(List<Document> documents) {
List<Map<String, Object>> metadataList = this.extract(documents);
for (int idx = 0; idx < documents.size(); idx++) {
documents.get(idx).getMetadata().putAll(metadataList.get(idx));
}
return documents;
}
/**
* Extracts metadata for a list of documents, returning a list of metadata
* dictionaries corresponding to each document.
* @param documents Documents to extract metadata from.
* @return List of metadata dictionaries corresponding to each document
*/
abstract public List<Map<String, Object>> extract(List<Document> documents);
}

View File

@@ -0,0 +1,92 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.loader.extractor;
import java.util.ArrayList;
import java.util.List;
import org.springframework.ai.document.ContentFormatter;
import org.springframework.ai.document.DefaultContentFormatter;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentTransformer;
/**
* @author Christian Tzolov
*/
public class ContentFormatEnricher implements DocumentTransformer {
/**
* Disable the content-formatter template rewrite.
*/
private boolean disableTemplateRewrite = false;
private ContentFormatter contentFormatter;
public ContentFormatEnricher(ContentFormatter contentFormatter) {
this(contentFormatter, false);
}
public ContentFormatEnricher(ContentFormatter contentFormatter, boolean disableTemplateRewrite) {
this.contentFormatter = contentFormatter;
this.disableTemplateRewrite = disableTemplateRewrite;
}
/**
* Post process documents chunked from loader. Allows extractors to be chained.
* @param documents to post process.
* @return
*/
public List<Document> apply(List<Document> documents) {
if (this.contentFormatter != null) {
documents.forEach(document -> {
// Update formatter
if (document.getContentFormatter() instanceof DefaultContentFormatter
&& this.contentFormatter instanceof DefaultContentFormatter) {
DefaultContentFormatter docFormatter = (DefaultContentFormatter) document.getContentFormatter();
DefaultContentFormatter toUpdateFormatter = (DefaultContentFormatter) this.contentFormatter;
var updatedEmbedExcludeKeys = new ArrayList<>(docFormatter.getExcludedEmbedMetadataKeys());
updatedEmbedExcludeKeys.addAll(toUpdateFormatter.getExcludedEmbedMetadataKeys());
var updatedInterfaceExcludeKeys = new ArrayList<>(docFormatter.getExcludedInferenceMetadataKeys());
updatedInterfaceExcludeKeys.addAll(toUpdateFormatter.getExcludedInferenceMetadataKeys());
var builder = DefaultContentFormatter.builder()
.withExcludedEmbedMetadataKeys(updatedEmbedExcludeKeys)
.withExcludedInferenceMetadataKeys(updatedInterfaceExcludeKeys)
.withMetadataTemplate(docFormatter.getMetadataTemplate())
.withMetadataSeparator(docFormatter.getMetadataSeparator());
if (!this.disableTemplateRewrite) {
builder.withTextTemplate(docFormatter.getTextTemplate());
}
document.setContentFormatter(builder.build());
}
else {
// Override formatter
document.setContentFormatter(this.contentFormatter);
}
});
}
return documents;
}
}

View File

@@ -0,0 +1,74 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.loader.extractor;
import java.util.List;
import java.util.Map;
import org.springframework.ai.client.AiClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentTransformer;
import org.springframework.ai.prompt.Prompt;
import org.springframework.ai.prompt.PromptTemplate;
import org.springframework.util.Assert;
/**
* Keyword extractor that uses model to extract 'excerpt_keywords' metadata field.
*
* @author Christian Tzolov
*/
public class KeywordExtractor implements DocumentTransformer {
private static final String EXCERPT_KEYWORDS_METADATA_KEY = "excerpt_keywords";
public static final String CONTEXT_STR_PLACEHOLDER = "context_str";
public static final String KEYWORDS_TEMPLATE = """
{context_str}. Give %s unique keywords for this
document. Format as comma separated. Keywords: """;
/**
* Model predictor
*/
private final AiClient aiClient;
/**
* The number of keywords to extract.
*/
private final int keywordCount;
public KeywordExtractor(AiClient aiClient, int keywordCount) {
Assert.notNull(aiClient, "AiClient must not be null");
Assert.isTrue(keywordCount >= 1, "Document count must be >= 1");
this.aiClient = aiClient;
this.keywordCount = keywordCount;
}
@Override
public List<Document> apply(List<Document> documents) {
for (Document document : documents) {
var template = new PromptTemplate(String.format(KEYWORDS_TEMPLATE, keywordCount));
Prompt prompt = template.create(Map.of(CONTEXT_STR_PLACEHOLDER, document.getContent()));
String keywords = this.aiClient.generate(prompt).getGeneration().getText();
document.getMetadata().putAll(Map.of(EXCERPT_KEYWORDS_METADATA_KEY, keywords));
}
return documents;
}
}

View File

@@ -0,0 +1,126 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.loader.extractor;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.springframework.ai.client.AiClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentTransformer;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.prompt.Prompt;
import org.springframework.ai.prompt.PromptTemplate;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
/**
* Title extractor with adjacent sharing that uses model to extract 'section_summary',
* 'prev_section_summary', 'next_section_summary' metadata fields.
*
* @author Christian Tzolov
*/
public class SummaryExtractor implements DocumentTransformer {
private static final String SECTION_SUMMARY_METADATA_KEY = "section_summary";
private static final String NEXT_SECTION_SUMMARY_METADATA_KEY = "next_section_summary";
private static final String PREV_SECTION_SUMMARY_METADATA_KEY = "prev_section_summary";
private static final String CONTEXT_STR_PLACEHOLDER = "context_str";
public static final String DEFAULT_SUMMARY_EXTRACT_TEMPLATE = """
Here is the content of the section:
{context_str}
Summarize the key topics and entities of the section.
Summary: """;
public enum SummaryType {
PREVIOUS, CURRENT, NEXT;
}
/**
* AI client.
*/
private final AiClient aiClient;
/**
* Number of documents from front to use for title extraction.
*/
private final List<SummaryType> summaryTypes;
private final MetadataMode metadataMode;
/**
* Template for summary extraction.
*/
private final String summaryTemplate;
public SummaryExtractor(AiClient aiClient, List<SummaryType> summaryTypes) {
this(aiClient, summaryTypes, DEFAULT_SUMMARY_EXTRACT_TEMPLATE, MetadataMode.ALL);
}
public SummaryExtractor(AiClient aiClient, List<SummaryType> summaryTypes, String summaryTemplate,
MetadataMode metadataMode) {
Assert.notNull(aiClient, "AiClient must not be null");
Assert.hasText(summaryTemplate, "Summary template must not be empty");
this.aiClient = aiClient;
this.summaryTypes = CollectionUtils.isEmpty(summaryTypes) ? List.of(SummaryType.CURRENT) : summaryTypes;
this.metadataMode = metadataMode;
this.summaryTemplate = summaryTemplate;
}
@Override
public List<Document> apply(List<Document> documents) {
List<String> documentSummaries = new ArrayList<>();
for (Document document : documents) {
var documentContext = document.getFormattedContent(this.metadataMode);
Prompt prompt = new PromptTemplate(this.summaryTemplate)
.create(Map.of(CONTEXT_STR_PLACEHOLDER, documentContext));
documentSummaries.add(this.aiClient.generate(prompt).getGeneration().getText());
}
for (int i = 0; i < documentSummaries.size(); i++) {
Map<String, Object> summaryMetadata = new HashMap<>();
if (i > 0 && this.summaryTypes.contains(SummaryType.PREVIOUS)) {
summaryMetadata.put(PREV_SECTION_SUMMARY_METADATA_KEY, documentSummaries.get(i - 1));
}
if (i < (documentSummaries.size() - 1) && this.summaryTypes.contains(SummaryType.NEXT)) {
summaryMetadata.put(NEXT_SECTION_SUMMARY_METADATA_KEY, documentSummaries.get(i + 1));
}
if (this.summaryTypes.contains(SummaryType.CURRENT)) {
summaryMetadata.put(SECTION_SUMMARY_METADATA_KEY, documentSummaries.get(i));
}
documents.get(i).getMetadata().putAll(summaryMetadata);
}
return documents;
}
}

View File

@@ -2,6 +2,8 @@ package org.springframework.ai.splitter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.ContentFormatter;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentTransformer;
@@ -15,24 +17,41 @@ public abstract class TextSplitter implements DocumentTransformer {
private static final Logger logger = LoggerFactory.getLogger(TextSplitter.class);
/**
* If true the children documents inherit the content-type of the parent they were
* split from.
*/
private boolean copyContentFormatter = true;
@Override
public List<Document> apply(List<Document> documents) {
return doSplitDocuments(documents);
}
public void setCopyContentFormatter(boolean copyContentFormatter) {
this.copyContentFormatter = copyContentFormatter;
}
public boolean isCopyContentFormatter() {
return this.copyContentFormatter;
}
private List<Document> doSplitDocuments(List<Document> documents) {
List<String> texts = new ArrayList<>();
Map<String, Object> metadata = new HashMap<>();
List<ContentFormatter> formatters = new ArrayList<>();
for (Document doc : documents) {
texts.add(doc.getText());
texts.add(doc.getContent());
metadata.putAll(doc.getMetadata());
formatters.add(doc.getContentFormatter());
}
return createDocuments(texts, metadata);
return createDocuments(texts, formatters, metadata);
}
private List<Document> createDocuments(List<String> texts, Map<String, Object> metadata) {
private List<Document> createDocuments(List<String> texts, List<ContentFormatter> formatters,
Map<String, Object> metadata) {
// Process the data in a column oriented way and recreate the Document
List<Document> documents = new ArrayList<>();
@@ -49,6 +68,13 @@ public abstract class TextSplitter implements DocumentTransformer {
.stream()
.collect(Collectors.toMap(e -> e.getKey(), e -> e.getValue()));
Document newDoc = new Document(chunk, metadataCopy);
if (this.copyContentFormatter) {
// Transfer the content-formatter of the parent to the chunked
// documents it was slit into.
newDoc.setContentFormatter(formatters.get(i));
}
// TODO copy over other properties.
documents.add(newDoc);
}

View File

@@ -1,10 +1,10 @@
package org.springframework.ai.vectorstore;
import org.springframework.ai.document.Document;
import java.util.List;
import java.util.Optional;
import org.springframework.ai.document.Document;
public interface VectorStore {
/**

View File

@@ -0,0 +1,94 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.document;
import java.util.Map;
import org.junit.jupiter.api.Test;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
*/
public class ContentFormatterTests {
Document document = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("embedKey1", "value1", "embedKey2", "value2", "embedKey3", "value3", "llmKey2", "value4"));
@Test
public void noExplicitlySetFormatter() {
assertThat(document.getContent()).isEqualTo("""
The World is Big and Salvation Lurks Around the Corner""");
assertThat(document.getFormattedContent()).isEqualTo(document.getFormattedContent(MetadataMode.ALL));
assertThat(document.getFormattedContent())
.isEqualTo(document.getFormattedContent(Document.DEFAULT_CONTENT_FORMATTER, MetadataMode.ALL));
}
@Test
public void defaultConfigTextFormatter() {
DefaultContentFormatter defaultConfigFormatter = DefaultContentFormatter.defaultConfig();
assertThat(document.getFormattedContent(defaultConfigFormatter, MetadataMode.ALL)).isEqualTo("""
llmKey2: value4
embedKey1: value1
embedKey2: value2
embedKey3: value3
The World is Big and Salvation Lurks Around the Corner""");
assertThat(document.getFormattedContent(defaultConfigFormatter, MetadataMode.ALL))
.isEqualTo(document.getFormattedContent());
assertThat(document.getFormattedContent(defaultConfigFormatter, MetadataMode.ALL))
.isEqualTo(defaultConfigFormatter.format(document, MetadataMode.ALL));
}
@Test
public void customTextFormatter() {
DefaultContentFormatter textFormatter = DefaultContentFormatter.builder()
.withExcludedEmbedMetadataKeys("embedKey2", "embedKey3")
.withExcludedInferenceMetadataKeys("llmKey2")
.withTextTemplate("Metadata:\n{metadata_string}\n\nText:{content}")
.withMetadataTemplate("Key/Value {key}={value}")
.build();
assertThat(document.getFormattedContent(textFormatter, MetadataMode.EMBED)).isEqualTo("""
Metadata:
Key/Value llmKey2=value4
Key/Value embedKey1=value1
Text:The World is Big and Salvation Lurks Around the Corner""");
assertThat(document.getContent()).isEqualTo("""
The World is Big and Salvation Lurks Around the Corner""");
assertThat(document.getFormattedContent(textFormatter, MetadataMode.EMBED))
.isEqualTo(textFormatter.format(document, MetadataMode.EMBED));
var documentWithCustomFormatter = new Document(document.getId(), document.getContent(), document.getMetadata());
documentWithCustomFormatter.setContentFormatter(textFormatter);
assertThat(document.getFormattedContent(textFormatter, MetadataMode.ALL))
.isEqualTo(documentWithCustomFormatter.getFormattedContent());
}
}

View File

@@ -40,7 +40,7 @@ public class JsonLoaderTests {
List<Document> documents = jsonLoader.load();
assertThat(documents).isNotEmpty();
for (Document document : documents) {
assertThat(document.getText()).isNotEmpty();
assertThat(document.getContent()).isNotEmpty();
}
}

View File

@@ -48,7 +48,7 @@ public class TextLoaderTests {
assertThat(document.getMetadata().get("customKey")).isEqualTo("Value");
assertThat(document.getMetadata().get(TextLoader.SOURCE_METADATA)).isEqualTo("text_source.txt");
assertThat(document.getMetadata().get(TextLoader.CHARSET_METADATA)).isEqualTo("UTF-8");
assertThat(document.getText()).isNotEmpty();
assertThat(document.getContent()).isNotEmpty();
}
}

View File

@@ -0,0 +1,108 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.splitter;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import static org.assertj.core.api.Assertions.assertThat;
import org.junit.jupiter.api.Test;
import org.springframework.ai.document.DefaultContentFormatter;
import org.springframework.ai.document.Document;
/**
* @author Christian Tzolov
*/
public class TextSplitterTests {
static TextSplitter testTextSplitter = new TextSplitter() {
@Override
protected List<String> splitText(String text) {
int chuckSize = text.length() / 2;
List<String> chunks = new ArrayList<>();
chunks.add(text.substring(0, chuckSize));
chunks.add(text.substring(chuckSize, text.length()));
return chunks;
}
};
@Test
public void testSplitText() {
var contentFormatter1 = DefaultContentFormatter.defaultConfig();
var contentFormatter2 = DefaultContentFormatter.defaultConfig();
assertThat(contentFormatter1).isNotSameAs(contentFormatter2);
var doc1 = new Document("In the end, writing arises when man realizes that memory is not enough.",
Map.of("key1", "value1", "key2", "value2"));
doc1.setContentFormatter(contentFormatter1);
var doc2 = new Document("The most oppressive thing about the labyrinth is that you are constantly "
+ "being forced to choose. It isnt the lack of an exit, but the abundance of exits that is so disorienting.",
Map.of("key2", "value22", "key3", "value3"));
doc2.setContentFormatter(contentFormatter2);
List<Document> chunks = testTextSplitter.apply(List.of(doc1, doc2));
assertThat(testTextSplitter.isCopyContentFormatter()).isTrue();
assertThat(chunks).hasSize(4);
// Doc1 chunks:
assertThat(chunks.get(0).getContent()).isEqualTo("In the end, writing arises when man");
assertThat(chunks.get(1).getContent()).isEqualTo(" realizes that memory is not enough.");
// Doc2 chunks:
assertThat(chunks.get(2).getContent())
.isEqualTo("The most oppressive thing about the labyrinth is that you are constantly being forced to ");
assertThat(chunks.get(3).getContent())
.isEqualTo("choose. It isnt the lack of an exit, but the abundance of exits that is so disorienting.");
// Verify that the same, merged metadata is copied to all chunks.
assertThat(chunks.get(0).getMetadata()).isEqualTo(chunks.get(1).getMetadata());
assertThat(chunks.get(0).getMetadata()).isEqualTo(chunks.get(2).getMetadata());
assertThat(chunks.get(0).getMetadata()).isEqualTo(chunks.get(3).getMetadata());
assertThat(chunks.get(0).getMetadata()).containsKeys("key1", "key2", "key3");
// Verify that the content formatters are copied from the parents to the chunks.
// doc1 -> chunk0, chunk1 and doc2 -> chunk2, chunk3
assertThat(chunks.get(0).getContentFormatter()).isSameAs(contentFormatter1);
assertThat(chunks.get(1).getContentFormatter()).isSameAs(contentFormatter1);
assertThat(chunks.get(2).getContentFormatter()).isSameAs(contentFormatter2);
assertThat(chunks.get(3).getContentFormatter()).isSameAs(contentFormatter2);
// Disable copy content formatters
testTextSplitter.setCopyContentFormatter(false);
chunks = testTextSplitter.apply(List.of(doc1, doc2));
assertThat(chunks.get(0).getContentFormatter()).isNotSameAs(contentFormatter1);
assertThat(chunks.get(1).getContentFormatter()).isNotSameAs(contentFormatter1);
assertThat(chunks.get(2).getContentFormatter()).isNotSameAs(contentFormatter2);
assertThat(chunks.get(3).getContentFormatter()).isNotSameAs(contentFormatter2);
}
}

View File

@@ -4,7 +4,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.2.0-SNAPSHOT</version>
<version>0.7.0-SNAPSHOT</version>
</parent>
<artifactId>spring-ai-docs</artifactId>
<name>Spring AI Docs</name>

View File

@@ -4,7 +4,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.2.0-SNAPSHOT</version>
<version>0.7.0-SNAPSHOT</version>
</parent>
<artifactId>spring-ai-openai</artifactId>
<packaging>jar</packaging>

View File

@@ -5,7 +5,6 @@ import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
import com.theokanning.openai.Usage;
import com.theokanning.openai.embedding.EmbeddingRequest;
@@ -14,6 +13,7 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.embedding.Embedding;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.embedding.EmbeddingResponse;
@@ -30,15 +30,23 @@ public class OpenAiEmbeddingClient implements EmbeddingClient {
private final AtomicInteger embeddingDimensions = new AtomicInteger(-1);
private final MetadataMode metadataMode;
public OpenAiEmbeddingClient(OpenAiService openAiService) {
this(openAiService, "text-embedding-ada-002");
}
public OpenAiEmbeddingClient(OpenAiService openAiService, String model) {
this(openAiService, model, MetadataMode.EMBED);
}
public OpenAiEmbeddingClient(OpenAiService openAiService, String model, MetadataMode metadataMode) {
Assert.notNull(openAiService, "OpenAiService must not be null");
Assert.notNull(model, "Model must not be null");
Assert.notNull(metadataMode, "metadataMode must not be null");
this.openAiService = openAiService;
this.model = model;
this.metadataMode = metadataMode;
}
@Override
@@ -51,7 +59,7 @@ public class OpenAiEmbeddingClient implements EmbeddingClient {
public List<Double> embed(Document document) {
EmbeddingRequest embeddingRequest = EmbeddingRequest.builder()
.input(List.of(document.getContent()))
.input(List.of(document.getFormattedContent(this.metadataMode)))
.model(this.model)
.build();
com.theokanning.openai.embedding.EmbeddingResult nativeEmbeddingResult = this.openAiService

View File

@@ -94,10 +94,6 @@ public class AcmeIT extends AbstractIT {
AiResponse response = aiClient.generate(prompt);
evaluateQuestionAndAnswer(userQuery, response, true);
// Chain
// qa = new ConversationalRetrievalChain(llmClient, userPromptTemplate,
// vectorStoreRetriever, )
}
private Message getSystemMessage(List<Document> similarDocuments) {

View File

@@ -0,0 +1,200 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.openai.extractor;
import java.io.IOException;
import java.time.Duration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import com.theokanning.openai.service.OpenAiService;
import org.junit.jupiter.api.Test;
import org.springframework.ai.document.DefaultContentFormatter;
import org.springframework.ai.document.Document;
import org.springframework.ai.loader.extractor.ContentFormatEnricher;
import org.springframework.ai.loader.extractor.KeywordExtractor;
import org.springframework.ai.loader.extractor.SummaryExtractor;
import org.springframework.ai.loader.extractor.SummaryExtractor.SummaryType;
import org.springframework.ai.openai.client.OpenAiClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
import org.springframework.util.StringUtils;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
*/
@SpringBootTest
public class MetadataExtractorIT {
@Autowired
KeywordExtractor keywordExtractor;
@Autowired
SummaryExtractor summaryExtractor;
@Autowired
ContentFormatEnricher metadataExtractor;
@Autowired
DefaultContentFormatter defaultContentFormatter;
Document document1 = new Document("Somewhere in the Andes, they believe to this very day that the"
+ " future is behind you. It comes up from behind your back, surprising and unforeseeable, while the past "
+ " is always before your eyes, that which has already happened. When they talk about the past, the people of"
+ " the Aymara tribe point in front of them. You walk forward facing the past and you turn back toward the future.",
new HashMap<>(Map.of("key", "value")));
Document document2 = new Document(
"The Spring Framework is divided into modules. Applications can choose which modules"
+ " they need. At the heart are the modules of the core container, including a configuration model and a "
+ "dependency injection mechanism. Beyond that, the Spring Framework provides foundational support "
+ " for different application architectures, including messaging, transactional data and persistence, "
+ "and web. It also includes the Servlet-based Spring MVC web framework and, in parallel, the Spring "
+ "WebFlux reactive web framework.");
@Test
public void testKeywordExtractor() {
var updatedDocuments = keywordExtractor.apply(List.of(document1, document2));
List<Map<String, Object>> keywords = updatedDocuments.stream().map(d -> d.getMetadata()).toList();
assertThat(updatedDocuments.size()).isEqualTo(2);
var keywords1 = keywords.get(0);
var keywords2 = keywords.get(1);
assertThat(keywords1).containsKeys("excerpt_keywords");
assertThat(keywords2).containsKeys("excerpt_keywords");
assertThat((String) keywords1.get("excerpt_keywords")).contains("Andes", "Aymara");
assertThat((String) keywords2.get("excerpt_keywords")).contains("Spring Framework", "dependency injection");
}
@Test
public void testSummaryExtractor() {
var updatedDocuments = summaryExtractor.apply(List.of(document1, document2));
List<Map<String, Object>> summaries = updatedDocuments.stream().map(d -> d.getMetadata()).toList();
assertThat(summaries.size()).isEqualTo(2);
var summary1 = summaries.get(0);
var summary2 = summaries.get(1);
assertThat(summary1).containsKeys("section_summary", "next_section_summary");
assertThat(summary1).doesNotContainKeys("prev_section_summary");
assertThat(summary2).containsKeys("section_summary", "prev_section_summary");
assertThat(summary2).doesNotContainKeys("next_section_summary");
assertThat((String) summary1.get("section_summary")).isNotEmpty();
assertThat((String) summary1.get("next_section_summary")).isNotEmpty();
assertThat((String) summary2.get("section_summary")).isNotEmpty();
assertThat((String) summary2.get("prev_section_summary")).isNotEmpty();
assertThat((String) summary1.get("section_summary")).isEqualTo((String) summary2.get("prev_section_summary"));
assertThat((String) summary1.get("next_section_summary")).isEqualTo((String) summary2.get("section_summary"));
}
@Test
public void testContentFormatEnricher() {
assertThat(((DefaultContentFormatter) document1.getContentFormatter()).getExcludedEmbedMetadataKeys())
.doesNotContain("NewEmbedKey");
assertThat(((DefaultContentFormatter) document1.getContentFormatter()).getExcludedInferenceMetadataKeys())
.doesNotContain("NewInferenceKey");
assertThat(((DefaultContentFormatter) document2.getContentFormatter()).getExcludedEmbedMetadataKeys())
.doesNotContain("NewEmbedKey");
assertThat(((DefaultContentFormatter) document2.getContentFormatter()).getExcludedInferenceMetadataKeys())
.doesNotContain("NewInferenceKey");
List<Document> enrichedDocuments = metadataExtractor.apply(List.of(document1, document2));
assertThat(enrichedDocuments.size()).isEqualTo(2);
var doc1 = enrichedDocuments.get(0);
var doc2 = enrichedDocuments.get(1);
assertThat(doc1).isEqualTo(document1);
assertThat(doc2).isEqualTo(document2);
assertThat(((DefaultContentFormatter) doc1.getContentFormatter()).getTextTemplate())
.isSameAs(defaultContentFormatter.getTextTemplate());
assertThat(((DefaultContentFormatter) doc1.getContentFormatter()).getExcludedEmbedMetadataKeys())
.contains("NewEmbedKey");
assertThat(((DefaultContentFormatter) doc1.getContentFormatter()).getExcludedInferenceMetadataKeys())
.contains("NewInferenceKey");
assertThat(((DefaultContentFormatter) doc2.getContentFormatter()).getTextTemplate())
.isSameAs(defaultContentFormatter.getTextTemplate());
assertThat(((DefaultContentFormatter) doc2.getContentFormatter()).getExcludedEmbedMetadataKeys())
.contains("NewEmbedKey");
assertThat(((DefaultContentFormatter) doc2.getContentFormatter()).getExcludedInferenceMetadataKeys())
.contains("NewInferenceKey");
}
@SpringBootConfiguration
public static class OpenAiTestConfiguration {
@Bean
public OpenAiService theoOpenAiService() throws IOException {
String apiKey = System.getenv("OPENAI_API_KEY");
if (!StringUtils.hasText(apiKey)) {
throw new IllegalArgumentException(
"You must provide an API key. Put it in an environment variable under the name OPENAI_API_KEY");
}
OpenAiService openAiService = new OpenAiService(apiKey, Duration.ofSeconds(60));
return openAiService;
}
@Bean
public OpenAiClient openAiClient(OpenAiService theoOpenAiService) {
OpenAiClient openAiClient = new OpenAiClient(theoOpenAiService);
openAiClient.setTemperature(0.3);
return openAiClient;
}
@Bean
public KeywordExtractor keywordExtractor(OpenAiClient aiClient) {
return new KeywordExtractor(aiClient, 5);
}
@Bean
public SummaryExtractor summaryExtractor(OpenAiClient aiClient) {
return new SummaryExtractor(aiClient, List.of(SummaryType.PREVIOUS, SummaryType.CURRENT, SummaryType.NEXT));
}
@Bean
public DefaultContentFormatter defaultContentFormatter() {
return DefaultContentFormatter.builder()
.withExcludedEmbedMetadataKeys("NewEmbedKey")
.withExcludedInferenceMetadataKeys("NewInferenceKey")
.build();
}
@Bean
public ContentFormatEnricher metadataExtractor(DefaultContentFormatter defaultContentFormatter) {
return new ContentFormatEnricher(defaultContentFormatter, false);
}
}
}

View File

@@ -23,7 +23,7 @@ import static org.assertj.core.api.Assertions.assertThat;
@SpringBootTest
public class SimplePersistentVectorStoreIT {
@Value("classpath:/data/acme/bikes.json")
@Value("file:src/test/resources/data/acme/bikes.json")
private Resource bikesJsonResource;
@Autowired

View File

@@ -4,7 +4,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.2.0-SNAPSHOT</version>
<version>0.7.0-SNAPSHOT</version>
</parent>
<artifactId>spring-ai-spring-boot-autoconfigure</artifactId>
<packaging>jar</packaging>

View File

@@ -4,7 +4,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.2.0-SNAPSHOT</version>
<version>0.7.0-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-azure-openai-spring-boot-starter</artifactId>

View File

@@ -4,7 +4,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.2.0-SNAPSHOT</version>
<version>0.7.0-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>

View File

@@ -5,7 +5,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.2.0-SNAPSHOT</version>
<version>0.7.0-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-milvus-store</artifactId>
@@ -22,7 +22,7 @@
<properties>
<milvus.version>2.3.0</milvus.version>
<spring-ai.version>0.2.0-SNAPSHOT</spring-ai.version>
<spring-ai.version>0.7.0-SNAPSHOT</spring-ai.version>
<!-- testing -->
<testcontainers.version>1.19.0</testcontainers.version>
<!-- <hikari-cp.version>4.0.3</hikari-cp.version> -->

View File

@@ -269,7 +269,7 @@ public class MilvusVectorStore implements VectorStore, SmartLifecycle {
docIdArray.add(document.getId());
// Use a (future) DocumentTextLayoutFormatter instance to extract
// the content used to compute the embeddings
contentArray.add(document.getText());
contentArray.add(document.getContent());
metadataArray.add(new JSONObject(document.getMetadata()));
embeddingArray.add(toFloatList(embedding));
}

View File

@@ -116,7 +116,7 @@ public class MilvusVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getText()).isEqualTo(
assertThat(resultDoc.getContent()).isEqualTo(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression");
assertThat(resultDoc.getMetadata()).hasSize(2);
assertThat(resultDoc.getMetadata()).containsKey("meta2");
@@ -152,7 +152,7 @@ public class MilvusVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getText()).isEqualTo("Spring AI rocks!!");
assertThat(resultDoc.getContent()).isEqualTo("Spring AI rocks!!");
assertThat(resultDoc.getMetadata()).containsKey("meta1");
assertThat(resultDoc.getMetadata()).containsKey("distance");
@@ -167,7 +167,7 @@ public class MilvusVectorStoreIT {
assertThat(results).hasSize(1);
resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getText()).isEqualTo("The World is Big and Salvation Lurks Around the Corner");
assertThat(resultDoc.getContent()).isEqualTo("The World is Big and Salvation Lurks Around the Corner");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
@@ -205,7 +205,7 @@ public class MilvusVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getText()).isEqualTo(
assertThat(resultDoc.getContent()).isEqualTo(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");

View File

@@ -5,7 +5,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.2.0-SNAPSHOT</version>
<version>0.7.0-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-neo4j-store</artifactId>
@@ -21,7 +21,7 @@
</scm>
<properties>
<spring-ai.version>0.2.0-SNAPSHOT</spring-ai.version>
<spring-ai.version>0.7.0-SNAPSHOT</spring-ai.version>
<!-- testing -->
<testcontainers.version>1.19.0</testcontainers.version>
</properties>

View File

@@ -325,7 +325,7 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean {
row.put("id", document.getId());
var properties = new HashMap<String, Object>();
properties.put("text", document.getText());
properties.put("text", document.getContent());
document.getMetadata().forEach((k, v) -> properties.put("metadata." + k, Values.value(v)));
row.put("properties", properties);

View File

@@ -70,7 +70,7 @@ class Neo4jVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId());
assertThat(resultDoc.getText()).isEqualTo(
assertThat(resultDoc.getContent()).isEqualTo(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
@@ -100,7 +100,7 @@ class Neo4jVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getText()).isEqualTo("Spring AI rocks!!");
assertThat(resultDoc.getContent()).isEqualTo("Spring AI rocks!!");
assertThat(resultDoc.getMetadata()).containsKey("meta1");
assertThat(resultDoc.getMetadata()).containsKey("distance");
@@ -115,7 +115,7 @@ class Neo4jVectorStoreIT {
assertThat(results).hasSize(1);
resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getText()).isEqualTo("The World is Big and Salvation Lurks Around the Corner");
assertThat(resultDoc.getContent()).isEqualTo("The World is Big and Salvation Lurks Around the Corner");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
@@ -144,7 +144,7 @@ class Neo4jVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId());
assertThat(resultDoc.getText()).isEqualTo(
assertThat(resultDoc.getContent()).isEqualTo(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");

View File

@@ -5,7 +5,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.2.0-SNAPSHOT</version>
<version>0.7.0-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-pgvector-store</artifactId>
@@ -23,7 +23,7 @@
<properties>
<pgvector.version>0.1.3</pgvector.version>
<postgresql.version>42.6.0</postgresql.version>
<spring-ai.version>0.2.0-SNAPSHOT</spring-ai.version>
<spring-ai.version>0.7.0-SNAPSHOT</spring-ai.version>
<!-- testing -->
<testcontainers.version>1.19.0</testcontainers.version>
<hikari-cp.version>4.0.3</hikari-cp.version>

View File

@@ -213,7 +213,7 @@ public class PgVectorStore implements VectorStore, SmartLifecycle {
document.setEmbedding(embedding);
UUID id = UUID.fromString(document.getId());
String content = document.getText();
String content = document.getContent();
Map<String, Object> metadata = document.getMetadata();
PGvector pgEmbedding = new PGvector(toFloatArray(embedding));

View File

@@ -96,7 +96,7 @@ public class PgVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getText()).isEqualTo(
assertThat(resultDoc.getContent()).isEqualTo(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression");
assertThat(resultDoc.getMetadata()).containsKeys("meta2", "distance");
@@ -129,7 +129,7 @@ public class PgVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getText()).isEqualTo("Spring AI rocks!!");
assertThat(resultDoc.getContent()).isEqualTo("Spring AI rocks!!");
assertThat(resultDoc.getMetadata()).containsKeys("meta1", "distance");
Document sameIdDocument = new Document(document.getId(),
@@ -143,7 +143,7 @@ public class PgVectorStoreIT {
assertThat(results).hasSize(1);
resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getText()).isEqualTo("The World is Big and Salvation Lurks Around the Corner");
assertThat(resultDoc.getContent()).isEqualTo("The World is Big and Salvation Lurks Around the Corner");
assertThat(resultDoc.getMetadata()).containsKeys("meta2", "distance");
});
}
@@ -179,7 +179,7 @@ public class PgVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getText()).isEqualTo(
assertThat(resultDoc.getContent()).isEqualTo(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression");
assertThat(resultDoc.getMetadata()).containsKeys("meta2", "distance");