Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add Chroma VectorStore support

Browse Source 
- Implement ChromaApi client, based on Chroma REST API.
 - Implement ChromaVectorStore, including support for filter expression conversion.
 - Common VectorStoreUtil class to share to/from Float/Double list/array convertion as well as Json/Map convertions.
 - Add ITs including for Basic Auth and Token autheticatios.
 - Add ChromaApi security support for BasicAuth and Token.
 - Fix an issue with Text filter expression parser, related to double-quoted identifiers.
 - Add Chroma README.md.
 - Add Chroma boot autoconfiguration

Resolves# #86
This commit is contained in:
Christian Tzolov
2023-11-08 18:16:14 +01:00
committed by Mark Pollack
parent 837c4080aa
commit 98ca3e2a8f
36 changed files with 2735 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
pom.xml
View File

@@ -32,6 +32,7 @@
<module>document-readers/tika-reader</module>
<module>embedding-clients/transformers-embedding</module>
<module>vector-stores/spring-ai-pinecone</module>
<module>vector-stores/spring-ai-chroma</module>
</modules>

7
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/filter/FilterExpressionTextParser.java
View File

@@ -191,11 +191,14 @@ public class FilterExpressionTextParser {
@Override
public Filter.Operand visitTextConstant(FiltersParser.TextConstantContext ctx) {
var twiceQuotedText = ctx.getText();
String onceQuotedText = twiceQuotedText.substring(1, twiceQuotedText.length() - 1);
String onceQuotedText = removeOuterQuotes(ctx.getText());
return new Filter.Value(onceQuotedText);
}
private String removeOuterQuotes(String in) {
return in.substring(1, in.length() - 1);
}
@Override
public Filter.Operand visitIntegerConstant(FiltersParser.IntegerConstantContext ctx) {
return new Filter.Value(Integer.valueOf(ctx.getText()));

27
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/filter/converter/AbstractFilterExpressionConverter.java
View File

@@ -19,6 +19,7 @@ package org.springframework.ai.vectorstore.filter.converter;
import java.util.List;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.Filter.Expression;
import org.springframework.ai.vectorstore.filter.Filter.ExpressionType;
import org.springframework.ai.vectorstore.filter.Filter.Group;
import org.springframework.ai.vectorstore.filter.Filter.Operand;
@@ -27,16 +28,20 @@ import org.springframework.util.Assert;
/**
* @author Christian Tzolov
*/
public abstract class AbstractFilterExpressionConverter {
public abstract class AbstractFilterExpressionConverter implements FilterExpressionConverter {
public String convert(Operand operand) {
Assert.notNull(operand, "Operand can't be null");
@Override
public String convertExpression(Expression expression) {
return this.convertOperand(expression);
}
protected String convertOperand(Operand operand) {
var context = new StringBuilder();
this.convert(operand, context);
this.convertOperand(operand, context);
return context.toString();
}
protected void convert(Operand operand, StringBuilder context) {
protected void convertOperand(Operand operand, StringBuilder context) {
if (operand instanceof Filter.Group group) {
this.doGroup(group, context);
@@ -88,7 +93,7 @@ public abstract class AbstractFilterExpressionConverter {
protected void doGroup(Group group, StringBuilder context) {
this.doStartGroup(group, context);
this.convert(group.content(), context);
this.convertOperand(group.content(), context);
this.doEndGroup(group, context);
}
@@ -110,4 +115,14 @@ public abstract class AbstractFilterExpressionConverter {
context.append(",");
}
// Utilities
protected boolean hasOuterQuotes(String str) {
str = str.trim();
return (str.startsWith("\"") && str.endsWith("\"")) || (str.startsWith("'") && str.endsWith("'"));
}
protected String removeOuterQuotes(String in) {
return in.substring(1, in.length() - 1);
}
}

29
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/filter/converter/ChromaFilterExpressionConverter.java Normal file
View File

@@ -0,0 +1,29 @@
/*
* 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.vectorstore.filter.converter;
import org.springframework.ai.vectorstore.filter.Filter;
/**
* Converts {@link Filter.Expression} into Chroma metadata filter expression format.
* (https://docs.trychroma.com/usage-guide#using-where-filters)
*
* @author Christian Tzolov
*/
public class ChromaFilterExpressionConverter extends PineconeFilterExpressionConverter {
}

32
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/filter/converter/FilterExpressionConverter.java Normal file
View File

@@ -0,0 +1,32 @@
/*
* 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.vectorstore.filter.converter;
import org.springframework.ai.vectorstore.filter.Filter;
/**
* Converters a generic, portable {@link Filter.Expression} into a
* {@link org.springframework.ai.vectorstore.VectorStore} specific expression language
* format.
*
* @author Christian Tzolov
*/
public interface FilterExpressionConverter {
public String convertExpression(Filter.Expression expression);
}

9
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/filter/converter/MilvusFilterExpressionConverter.java
View File

@@ -31,9 +31,9 @@ public class MilvusFilterExpressionConverter extends AbstractFilterExpressionCon
@Override
protected void doExpression(Expression exp, StringBuilder context) {
this.convert(exp.left(), context);
this.convertOperand(exp.left(), context);
context.append(getOperationSymbol(exp));
this.convert(exp.right(), context);
this.convertOperand(exp.right(), context);
}
private String getOperationSymbol(Expression exp) {
@@ -65,12 +65,13 @@ public class MilvusFilterExpressionConverter extends AbstractFilterExpressionCon
@Override
protected void doGroup(Group group, StringBuilder context) {
this.convert(new Expression(ExpressionType.AND, group.content(), group.content()), context); // trick
this.convertOperand(new Expression(ExpressionType.AND, group.content(), group.content()), context); // trick
}
@Override
protected void doKey(Key key, StringBuilder context) {
context.append("metadata[\"" + key.key() + "\"]");
var identifier = (hasOuterQuotes(key.key())) ? removeOuterQuotes(key.key()) : key.key();
context.append("metadata[\"" + identifier + "\"]");
}
}

4
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/filter/converter/PgVectorFilterExpressionConverter.java
View File

@@ -30,9 +30,9 @@ public class PgVectorFilterExpressionConverter extends AbstractFilterExpressionC
@Override
protected void doExpression(Expression expression, StringBuilder context) {
this.convert(expression.left(), context);
this.convertOperand(expression.left(), context);
context.append(getOperationSymbol(expression));
this.convert(expression.right(), context);
this.convertOperand(expression.right(), context);
}
private String getOperationSymbol(Expression exp) {

11
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/filter/converter/PineconeFilterExpressionConverter.java
View File

@@ -35,16 +35,16 @@ public class PineconeFilterExpressionConverter extends AbstractFilterExpressionC
if (exp.type() == ExpressionType.AND || exp.type() == ExpressionType.OR) {
context.append(getOperationSymbol(exp));
context.append("[");
this.convert(exp.left(), context);
this.convertOperand(exp.left(), context);
context.append(",");
this.convert(exp.right(), context);
this.convertOperand(exp.right(), context);
context.append("]");
}
else {
this.convert(exp.left(), context);
this.convertOperand(exp.left(), context);
context.append("{");
context.append(getOperationSymbol(exp));
this.convert(exp.right(), context);
this.convertOperand(exp.right(), context);
context.append("}");
}
context.append("}");
@@ -57,7 +57,8 @@ public class PineconeFilterExpressionConverter extends AbstractFilterExpressionC
@Override
protected void doKey(Key key, StringBuilder context) {
context.append("\"" + key.key() + "\": ");
var identifier = (hasOuterQuotes(key.key())) ? removeOuterQuotes(key.key()) : key.key();
context.append("\"" + identifier + "\": ");
}
}

4
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/filter/converter/PrintFilterExpressionConverter.java
View File

@@ -28,9 +28,9 @@ import org.springframework.ai.vectorstore.filter.Filter.Key;
public class PrintFilterExpressionConverter extends AbstractFilterExpressionConverter {
public void doExpression(Expression expression, StringBuilder context) {
this.convert(expression.left(), context);
this.convertOperand(expression.left(), context);
context.append(" " + expression.type() + " ");
this.convert(expression.right(), context);
this.convertOperand(expression.right(), context);
}

2
spring-ai-core/src/test/java/org/springframework/ai/vectorstore/filter/FilterExpressionBuilderTests.java
View File

@@ -18,7 +18,6 @@ package org.springframework.ai.vectorstore.filter;
import java.util.List;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.springframework.ai.vectorstore.filter.Filter.Expression;
@@ -44,7 +43,6 @@ public class FilterExpressionBuilderTests {
@Test
public void testEQ() {
Expression expression = b.eq("country", "BG").build();
// country == "BG"
assertThat(b.eq("country", "BG").build()).isEqualTo(new Expression(EQ, new Key("country"), new Value("BG")));
}

12
spring-ai-core/src/test/java/org/springframework/ai/vectorstore/filter/FilterExpressionTextParserTests.java
View File

@@ -123,4 +123,16 @@ public class FilterExpressionTextParserTests {
assertThat(parser.getCache().get("WHERE " + expText)).isEqualTo(exp);
}
@Test
public void testIdentifiers() {
Expression exp = parser.parse("'country.1' == 'BG'");
assertThat(exp).isEqualTo(new Expression(EQ, new Key("'country.1'"), new Value("BG")));
exp = parser.parse("'country_1_2_3' == 'BG'");
assertThat(exp).isEqualTo(new Expression(EQ, new Key("'country_1_2_3'"), new Value("BG")));
exp = parser.parse("\"country 1 2 3\" == 'BG'");
assertThat(exp).isEqualTo(new Expression(EQ, new Key("\"country 1 2 3\""), new Value("BG")));
}
}

33
spring-ai-core/src/test/java/org/springframework/ai/vectorstore/filter/converter/MilvusFilterExpressionConverterTests.java
View File

@@ -40,12 +40,12 @@ import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.OR
*/
public class MilvusFilterExpressionConverterTests {
MilvusFilterExpressionConverter converter = new MilvusFilterExpressionConverter();
FilterExpressionConverter converter = new MilvusFilterExpressionConverter();
@Test
public void testEQ() {
// country == "BG"
String vectorExpr = converter.convert(new Expression(EQ, new Key("country"), new Value("BG")));
String vectorExpr = converter.convertExpression(new Expression(EQ, new Key("country"), new Value("BG")));
assertThat(vectorExpr).isEqualTo("metadata[\"country\"] == \"BG\"");
}
@@ -53,7 +53,7 @@ public class MilvusFilterExpressionConverterTests {
public void tesEqAndGte() {
// genre == "drama" AND year >= 2020
String vectorExpr = converter
.convert(new Expression(AND, new Expression(EQ, new Key("genre"), new Value("drama")),
.convertExpression(new Expression(AND, new Expression(EQ, new Key("genre"), new Value("drama")),
new Expression(GTE, new Key("year"), new Value(2020))));
assertThat(vectorExpr).isEqualTo("metadata[\"genre\"] == \"drama\" && metadata[\"year\"] >= 2020");
}
@@ -61,17 +61,18 @@ public class MilvusFilterExpressionConverterTests {
@Test
public void tesIn() {
// genre in ["comedy", "documentary", "drama"]
String vectorExpr = converter
.convert(new Expression(IN, new Key("genre"), new Value(List.of("comedy", "documentary", "drama"))));
String vectorExpr = converter.convertExpression(
new Expression(IN, new Key("genre"), new Value(List.of("comedy", "documentary", "drama"))));
assertThat(vectorExpr).isEqualTo("metadata[\"genre\"] in [\"comedy\",\"documentary\",\"drama\"]");
}
@Test
public void testNe() {
// year >= 2020 OR country == "BG" AND city != "Sofia"
String vectorExpr = converter.convert(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(AND, new Expression(EQ, new Key("country"), new Value("BG")),
new Expression(NE, new Key("city"), new Value("Sofia")))));
String vectorExpr = converter
.convertExpression(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(AND, new Expression(EQ, new Key("country"), new Value("BG")),
new Expression(NE, new Key("city"), new Value("Sofia")))));
assertThat(vectorExpr).isEqualTo(
"metadata[\"year\"] >= 2020 || metadata[\"country\"] == \"BG\" && metadata[\"city\"] != \"Sofia\"");
}
@@ -79,7 +80,7 @@ public class MilvusFilterExpressionConverterTests {
@Test
public void testGroup() {
// (year >= 2020 OR country == "BG") AND city NIN ["Sofia", "Plovdiv"]
String vectorExpr = converter.convert(new Expression(AND,
String vectorExpr = converter.convertExpression(new Expression(AND,
new Group(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(EQ, new Key("country"), new Value("BG")))),
new Expression(NIN, new Key("city"), new Value(List.of("Sofia", "Plovdiv")))));
@@ -90,7 +91,7 @@ public class MilvusFilterExpressionConverterTests {
@Test
public void tesBoolean() {
// isOpen == true AND year >= 2020 AND country IN ["BG", "NL", "US"]
String vectorExpr = converter.convert(new Expression(AND,
String vectorExpr = converter.convertExpression(new Expression(AND,
new Expression(AND, new Expression(EQ, new Key("isOpen"), new Value(true)),
new Expression(GTE, new Key("year"), new Value(2020))),
new Expression(IN, new Key("country"), new Value(List.of("BG", "NL", "US")))));
@@ -103,10 +104,20 @@ public class MilvusFilterExpressionConverterTests {
public void testDecimal() {
// temperature >= -15.6 && temperature <= +20.13
String vectorExpr = converter
.convert(new Expression(AND, new Expression(GTE, new Key("temperature"), new Value(-15.6)),
.convertExpression(new Expression(AND, new Expression(GTE, new Key("temperature"), new Value(-15.6)),
new Expression(LTE, new Key("temperature"), new Value(20.13))));
assertThat(vectorExpr).isEqualTo("metadata[\"temperature\"] >= -15.6 && metadata[\"temperature\"] <= 20.13");
}
@Test
public void testComplexIdentifiers() {
String vectorExpr = converter
.convertExpression(new Expression(EQ, new Key("\"country 1 2 3\""), new Value("BG")));
assertThat(vectorExpr).isEqualTo("metadata[\"country 1 2 3\"] == \"BG\"");
vectorExpr = converter.convertExpression(new Expression(EQ, new Key("'country 1 2 3'"), new Value("BG")));
assertThat(vectorExpr).isEqualTo("metadata[\"country 1 2 3\"] == \"BG\"");
}
}

30
spring-ai-core/src/test/java/org/springframework/ai/vectorstore/filter/converter/PgVectorFilterExpressionConverterTests.java
View File

@@ -40,12 +40,12 @@ import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.OR
*/
public class PgVectorFilterExpressionConverterTests {
PgVectorFilterExpressionConverter converter = new PgVectorFilterExpressionConverter();
FilterExpressionConverter converter = new PgVectorFilterExpressionConverter();
@Test
public void testEQ() {
// country == "BG"
String vectorExpr = converter.convert(new Expression(EQ, new Key("country"), new Value("BG")));
String vectorExpr = converter.convertExpression(new Expression(EQ, new Key("country"), new Value("BG")));
assertThat(vectorExpr).isEqualTo("$.country == \"BG\"");
}
@@ -53,7 +53,7 @@ public class PgVectorFilterExpressionConverterTests {
public void tesEqAndGte() {
// genre == "drama" AND year >= 2020
String vectorExpr = converter
.convert(new Expression(AND, new Expression(EQ, new Key("genre"), new Value("drama")),
.convertExpression(new Expression(AND, new Expression(EQ, new Key("genre"), new Value("drama")),
new Expression(GTE, new Key("year"), new Value(2020))));
assertThat(vectorExpr).isEqualTo("$.genre == \"drama\" && $.year >= 2020");
}
@@ -61,24 +61,25 @@ public class PgVectorFilterExpressionConverterTests {
@Test
public void tesIn() {
// genre in ["comedy", "documentary", "drama"]
String vectorExpr = converter
.convert(new Expression(IN, new Key("genre"), new Value(List.of("comedy", "documentary", "drama"))));
String vectorExpr = converter.convertExpression(
new Expression(IN, new Key("genre"), new Value(List.of("comedy", "documentary", "drama"))));
assertThat(vectorExpr).isEqualTo("$.genre in [\"comedy\",\"documentary\",\"drama\"]");
}
@Test
public void testNe() {
// year >= 2020 OR country == "BG" AND city != "Sofia"
String vectorExpr = converter.convert(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(AND, new Expression(EQ, new Key("country"), new Value("BG")),
new Expression(NE, new Key("city"), new Value("Sofia")))));
String vectorExpr = converter
.convertExpression(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(AND, new Expression(EQ, new Key("country"), new Value("BG")),
new Expression(NE, new Key("city"), new Value("Sofia")))));
assertThat(vectorExpr).isEqualTo("$.year >= 2020 || $.country == \"BG\" && $.city != \"Sofia\"");
}
@Test
public void testGroup() {
// (year >= 2020 OR country == "BG") AND city NIN ["Sofia", "Plovdiv"]
String vectorExpr = converter.convert(new Expression(AND,
String vectorExpr = converter.convertExpression(new Expression(AND,
new Group(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(EQ, new Key("country"), new Value("BG")))),
new Expression(NIN, new Key("city"), new Value(List.of("Sofia", "Plovdiv")))));
@@ -89,7 +90,7 @@ public class PgVectorFilterExpressionConverterTests {
@Test
public void tesBoolean() {
// isOpen == true AND year >= 2020 AND country IN ["BG", "NL", "US"]
String vectorExpr = converter.convert(new Expression(AND,
String vectorExpr = converter.convertExpression(new Expression(AND,
new Expression(AND, new Expression(EQ, new Key("isOpen"), new Value(true)),
new Expression(GTE, new Key("year"), new Value(2020))),
new Expression(IN, new Key("country"), new Value(List.of("BG", "NL", "US")))));
@@ -101,10 +102,17 @@ public class PgVectorFilterExpressionConverterTests {
public void testDecimal() {
// temperature >= -15.6 && temperature <= +20.13
String vectorExpr = converter
.convert(new Expression(AND, new Expression(GTE, new Key("temperature"), new Value(-15.6)),
.convertExpression(new Expression(AND, new Expression(GTE, new Key("temperature"), new Value(-15.6)),
new Expression(LTE, new Key("temperature"), new Value(20.13))));
assertThat(vectorExpr).isEqualTo("$.temperature >= -15.6 && $.temperature <= 20.13");
}
@Test
public void testComplexIdentifiers() {
String vectorExpr = converter
.convertExpression(new Expression(EQ, new Key("\"country 1 2 3\""), new Value("BG")));
assertThat(vectorExpr).isEqualTo("$.\"country 1 2 3\" == \"BG\"");
}
}

33
spring-ai-core/src/test/java/org/springframework/ai/vectorstore/filter/converter/PineconeFilterExpressionConverterTests.java
View File

@@ -40,12 +40,12 @@ import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.OR
*/
public class PineconeFilterExpressionConverterTests {
PineconeFilterExpressionConverter converter = new PineconeFilterExpressionConverter();
FilterExpressionConverter converter = new PineconeFilterExpressionConverter();
@Test
public void testEQ() {
// country == "BG"
String vectorExpr = converter.convert(new Expression(EQ, new Key("country"), new Value("BG")));
String vectorExpr = converter.convertExpression(new Expression(EQ, new Key("country"), new Value("BG")));
assertThat(vectorExpr).isEqualTo("{\"country\": {\"$eq\": \"BG\"}}");
}
@@ -53,7 +53,7 @@ public class PineconeFilterExpressionConverterTests {
public void tesEqAndGte() {
// genre == "drama" AND year >= 2020
String vectorExpr = converter
.convert(new Expression(AND, new Expression(EQ, new Key("genre"), new Value("drama")),
.convertExpression(new Expression(AND, new Expression(EQ, new Key("genre"), new Value("drama")),
new Expression(GTE, new Key("year"), new Value(2020))));
assertThat(vectorExpr)
.isEqualTo("{\"$and\": [{\"genre\": {\"$eq\": \"drama\"}},{\"year\": {\"$gte\": 2020}}]}");
@@ -62,17 +62,18 @@ public class PineconeFilterExpressionConverterTests {
@Test
public void tesIn() {
// genre in ["comedy", "documentary", "drama"]
String vectorExpr = converter
.convert(new Expression(IN, new Key("genre"), new Value(List.of("comedy", "documentary", "drama"))));
String vectorExpr = converter.convertExpression(
new Expression(IN, new Key("genre"), new Value(List.of("comedy", "documentary", "drama"))));
assertThat(vectorExpr).isEqualTo("{\"genre\": {\"$in\": [\"comedy\",\"documentary\",\"drama\"]}}");
}
@Test
public void testNe() {
// year >= 2020 OR country == "BG" AND city != "Sofia"
String vectorExpr = converter.convert(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(AND, new Expression(EQ, new Key("country"), new Value("BG")),
new Expression(NE, new Key("city"), new Value("Sofia")))));
String vectorExpr = converter
.convertExpression(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(AND, new Expression(EQ, new Key("country"), new Value("BG")),
new Expression(NE, new Key("city"), new Value("Sofia")))));
assertThat(vectorExpr).isEqualTo(
"{\"$or\": [{\"year\": {\"$gte\": 2020}},{\"$and\": [{\"country\": {\"$eq\": \"BG\"}},{\"city\": {\"$ne\": \"Sofia\"}}]}]}");
}
@@ -80,7 +81,7 @@ public class PineconeFilterExpressionConverterTests {
@Test
public void testGroup() {
// (year >= 2020 OR country == "BG") AND city NIN ["Sofia", "Plovdiv"]
String vectorExpr = converter.convert(new Expression(AND,
String vectorExpr = converter.convertExpression(new Expression(AND,
new Group(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(EQ, new Key("country"), new Value("BG")))),
new Expression(NIN, new Key("city"), new Value(List.of("Sofia", "Plovdiv")))));
@@ -91,7 +92,7 @@ public class PineconeFilterExpressionConverterTests {
@Test
public void tesBoolean() {
// isOpen == true AND year >= 2020 AND country IN ["BG", "NL", "US"]
String vectorExpr = converter.convert(new Expression(AND,
String vectorExpr = converter.convertExpression(new Expression(AND,
new Expression(AND, new Expression(EQ, new Key("isOpen"), new Value(true)),
new Expression(GTE, new Key("year"), new Value(2020))),
new Expression(IN, new Key("country"), new Value(List.of("BG", "NL", "US")))));
@@ -104,11 +105,21 @@ public class PineconeFilterExpressionConverterTests {
public void testDecimal() {
// temperature >= -15.6 && temperature <= +20.13
String vectorExpr = converter
.convert(new Expression(AND, new Expression(GTE, new Key("temperature"), new Value(-15.6)),
.convertExpression(new Expression(AND, new Expression(GTE, new Key("temperature"), new Value(-15.6)),
new Expression(LTE, new Key("temperature"), new Value(20.13))));
assertThat(vectorExpr)
.isEqualTo("{\"$and\": [{\"temperature\": {\"$gte\": -15.6}},{\"temperature\": {\"$lte\": 20.13}}]}");
}
@Test
public void testComplexIdentifiers() {
String vectorExpr = converter
.convertExpression(new Expression(EQ, new Key("\"country 1 2 3\""), new Value("BG")));
assertThat(vectorExpr).isEqualTo("{\"country 1 2 3\": {\"$eq\": \"BG\"}}");
vectorExpr = converter.convertExpression(new Expression(EQ, new Key("'country 1 2 3'"), new Value("BG")));
assertThat(vectorExpr).isEqualTo("{\"country 1 2 3\": {\"$eq\": \"BG\"}}");
}
}

8
spring-ai-spring-boot-autoconfigure/pom.xml
View File

@@ -94,6 +94,14 @@
<optional>true</optional>
</dependency>
<!-- Chroma Vector Store -->
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-chroma-store</artifactId>
<version>${project.parent.version}</version>
<optional>true</optional>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>

79
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaApiProperties.java Normal file
View File

@@ -0,0 +1,79 @@
/*
* 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.autoconfigure.vectorstore.chroma;
import org.springframework.boot.context.properties.ConfigurationProperties;
/**
* @author Christian Tzolov
*/
@ConfigurationProperties(ChromaApiProperties.CONFIG_PREFIX)
public class ChromaApiProperties {
public static final String CONFIG_PREFIX = "spring.ai.vectorstore.chroma.client";
private String host = "http://localhost";
private int port = 8000;
private String keyToken;
private String username;
private String password;
public String getHost() {
return host;
}
public void setHost(String baseUrl) {
this.host = baseUrl;
}
public int getPort() {
return port;
}
public void setPort(int port) {
this.port = port;
}
public String getKeyToken() {
return keyToken;
}
public void setKeyToken(String keyToken) {
this.keyToken = keyToken;
}
public String getUsername() {
return username;
}
public void setUsername(String username) {
this.username = username;
}
public String getPassword() {
return password;
}
public void setPassword(String password) {
this.password = password;
}
}

72
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreAutoConfiguration.java Normal file
View File

@@ -0,0 +1,72 @@
/*
* 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.autoconfigure.vectorstore.chroma;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.experimental.ai.chroma.ChromaApi;
import org.springframework.experimental.ai.vectorsore.ChromaVectorStore;
import org.springframework.util.StringUtils;
import org.springframework.web.client.RestTemplate;
/**
* @author Christian Tzolov
*/
@AutoConfiguration
@ConditionalOnClass({ EmbeddingClient.class, RestTemplate.class, ChromaVectorStore.class, ObjectMapper.class })
@EnableConfigurationProperties({ ChromaApiProperties.class, ChromaVectorStoreProperties.class })
public class ChromaVectorStoreAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public RestTemplate restTemplate() {
return new RestTemplate();
}
@Bean
@ConditionalOnMissingBean
public ChromaApi chromaApi(ChromaApiProperties apiProperties, RestTemplate restTemplate) {
String chromaUrl = String.format("%s:%s", apiProperties.getHost(), apiProperties.getPort());
var chromaApi = new ChromaApi(chromaUrl, restTemplate, new ObjectMapper());
if (StringUtils.hasText(apiProperties.getKeyToken())) {
chromaApi.withKeyToken(apiProperties.getKeyToken());
}
else if (StringUtils.hasText(apiProperties.getUsername()) && StringUtils.hasText(apiProperties.getPassword())) {
chromaApi.withBasicAuthCredentials(apiProperties.getUsername(), apiProperties.getPassword());
}
return chromaApi;
}
@Bean
@ConditionalOnMissingBean
public VectorStore vectorStore(EmbeddingClient embeddingClient, ChromaApi chromaApi,
ChromaVectorStoreProperties storeProperties) {
return new ChromaVectorStore(embeddingClient, chromaApi, storeProperties.getCollectionName());
}
}

40
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreProperties.java Normal file
View File

@@ -0,0 +1,40 @@
/*
* 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.autoconfigure.vectorstore.chroma;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.experimental.ai.vectorsore.ChromaVectorStore;
/**
* @author Christian Tzolov
*/
@ConfigurationProperties(ChromaVectorStoreProperties.CONFIG_PREFIX)
public class ChromaVectorStoreProperties {
public static final String CONFIG_PREFIX = "spring.ai.vectorstore.chroma.store";
private String collectionName = ChromaVectorStore.DEFAULT_COLLECTION_NAME;
public String getCollectionName() {
return collectionName;
}
public void setCollectionName(String collectionName) {
this.collectionName = collectionName;
}
}

1
spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports
View File

@@ -5,3 +5,4 @@ org.springframework.ai.autoconfigure.vectorstore.pinecone.PineconeVectorStoreAut
org.springframework.ai.autoconfigure.vectorstore.milvus.MilvusVectorStoreAutoConfiguration
org.springframework.ai.autoconfigure.embedding.transformer.TransformersEmbeddingClientAutoConfiguration
org.springframework.ai.autoconfigure.huggingface.HuggingfaceAutoConfiguration
org.springframework.ai.autoconfigure.vectorstore.chroma.ChromaVectorStoreAutoConfiguration

100
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreAutoConfigurationIT.java Normal file
View File

@@ -0,0 +1,100 @@
/*
* 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.autoconfigure.vectorstore.chroma;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.embedding.TransformersEmbeddingClient;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
*/
@Testcontainers
public class ChromaVectorStoreAutoConfigurationIT {
@Container
static GenericContainer<?> chromaContainer = new GenericContainer<>("ghcr.io/chroma-core/chroma:0.4.15")
.withExposedPorts(8000);
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(ChromaVectorStoreAutoConfiguration.class))
.withUserConfiguration(Config.class)
.withPropertyValues("spring.ai.vectorstore.chroma.client.host=http://localhost",
"spring.ai.vectorstore.chroma.client.port=" + chromaContainer.getMappedPort(8000),
"spring.ai.vectorstore.chroma.store.collectionName=TestCollection");
@Test
public void addAndSearchWithFilters() {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "Bulgaria"));
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "Netherland"));
vectorStore.add(List.of(bgDocument, nlDocument));
var request = SearchRequest.query("The World").withTopK(5);
List<Document> results = vectorStore.similaritySearch(request);
assertThat(results).hasSize(2);
results = vectorStore
.similaritySearch(request.withSimilarityThresholdAll().withFilterExpression("country == 'Bulgaria'"));
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId());
results = vectorStore
.similaritySearch(request.withSimilarityThresholdAll().withFilterExpression("country == 'Netherland'"));
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId());
// Remove all documents from the store
vectorStore.delete(List.of(bgDocument, nlDocument).stream().map(doc -> doc.getId()).toList());
});
}
@Configuration(proxyBeanMethods = false)
static class Config {
@Bean
public EmbeddingClient embeddingClient() {
return new TransformersEmbeddingClient();
}
}
}

3
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfigurationIT.java
View File

@@ -22,12 +22,10 @@ import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.ResourceUtils;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.embedding.TransformersEmbeddingClient;
@@ -39,6 +37,7 @@ import org.springframework.boot.autoconfigure.jdbc.JdbcTemplateAutoConfiguration
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.DefaultResourceLoader;
import static org.assertj.core.api.Assertions.assertThat;

154
vector-stores/spring-ai-chroma/README.md Normal file
View File

@@ -0,0 +1,154 @@
# Chroma VectorStore
This readme will walk you through setting up the Chroma VectorStore to store document embeddings and perform similarity searches.
<https://github.com/chroma-core/chroma/pkgs/container/chroma>
## What is Chroma?
[Chroma](https://docs.trychroma.com/) is the open-source embedding database. It gives you the tools to store document embeddings, content and metadata and to search through those embeddings including metadata filtering.
## Prerequisites
1. OpenAI Account: Create an account at [OpenAI Signup](https://platform.openai.com/signup) and generate the token at [API Keys](https://platform.openai.com/account/api-keys).
2. Access to ChromeDB. The [setup local ChromaDB](#appendix_a) appendix show how to setup a DB locally with a Docker container.
On startup the `ChromaVectorStore` creates the required collection if one is not provisioned already.
## Configuration
To set up ChromaVectorStore, you'll need to provide your OpenAI API Key. Set it as an environment variable like so:
```bash
export SPRING_AI_OPENAI_API_KEY='Your_OpenAI_API_Key'
```
## Dependencies
Add these dependencies to your project:
1. OpenAI: Required for calculating embeddings.
```xml
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
<version>0.7.0-SNAPSHOT</version>
</dependency>
```
2. Chroma VectorStore.
```xml
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-chroma-store</artifactId>
<version>0.7.0-SNAPSHOT</version>
</dependency>
```
## Sample Code
Create an `RestTemplate` instance with proper ChromaDB authorization configurations and Use it to create `ChromaApi` instance:
```java
@Bean
public RestTemplate restTemplate() {
return new RestTemplate();
}
@Bean
public ChromaApi chromaApi(RestTemplate restTemplate) {
String chromaUrl = "http://localhost:8000";
ChromaApi chromaApi = ChromaApi(chromaUrl, restTemplate);
return chromaApi;
}
```
> [!NOTE]
> For ChromaDB secured with [Static API Token Authentication](https://docs.trychroma.com/usage-guide#static-api-token-authentication) use the `ChromaApi#withKeyToken(<Your Token Credentials>)` method to set your credentials. Check the `ChromaWhereIT` for an example.
> [!NOTE]
> For ChromaDB secured with [Basic Authentication](https://docs.trychroma.com/usage-guide#basic-authentication) use the `ChromaApi#withBasicAuth(<your user>, <your password>)` method to set your credentials. Check the `BasicAuthChromaWhereIT` for an example.
Integrate with OpenAI's embeddings by adding the Spring Boot OpenAI starter to your project.
This provides you with an implementation of the Embeddings client:
```java
@Bean
public VectorStore chromaVectorStore(EmbeddingClient embeddingClient, ChromaApi chromaApi) {
return new ChromaVectorStore(embeddingClient, chromaApi, "TestCollection");
}
```
In your main code, create some documents
```java
List<Document> documents = List.of(
new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", Map.of("meta1", "meta1")),
new Document("The World is Big and Salvation Lurks Around the Corner"),
new Document("You walk forward facing the past and you turn back toward the future.", Map.of("meta2", "meta2")));
```
Add the documents to your vector store:
```java
vectorStore.add(List.of(document));
```
And finally, retrieve documents similar to a query:
```java
List<Document> results = vectorStore.similaritySearch("Spring", 5);
```
If all goes well, you should retrieve the document containing the text "Spring AI rocks!!".
### Metadata filtering
You can leverage the generic, portable [metadata filters](https://docs.spring.io/spring-ai/reference/api/vectordbs.html#_metadata_filters) with ChromaVector store as well.
For example you can use either the text expression language:
```java
vectorStore.similaritySearch("The World", TOP_K, SIMILARITY_THRESHOLD,
"author in ['john', 'jill'] && article_type == 'blog'");
```
or programmatically using the `Filter.Expression` DSL:
```java
FilterExpressionBuilder b = new FilterExpressionBuilder();
vectorStore.similaritySearch("The World", TOP_K, SIMILARITY_THRESHOLD,
b.and(
b.in(List.of("john", "jill")),
b.eq("article_type", "blog")).build());
```
NOTE: Those (portable) filter expressions get automatically converted into the proprietary Chroma `where` [filter expressions](https://docs.trychroma.com/usage-guide#using-where-filters).
For example this portable filter expression:
```sql
author in ['john', 'jill'] && article_type == 'blog'
```
is converted into the proprietary Chroma format:
```json
{"$and":[
{"author": {"$in": ["john", "jill"]}},
{"article_type":{"$eq":"blog"}}]
}"
```
## <a name="appendix_a" /> Appendix A: Run Chroma Locally
```
docker run -it --rm --name chroma -p 8000:8000 ghcr.io/chroma-core/chroma:0.4.15
```
starts a chroma store at <http://localhost:8000/api/v1>

85
vector-stores/spring-ai-chroma/pom.xml Normal file
View File

@@ -0,0 +1,85 @@
<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
<modelVersion>4.0.0</modelVersion>
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.1-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-chroma-store</artifactId>
<packaging>jar</packaging>
<name>Spring AI Chroma Vector Store</name>
<description>Spring AI Chroma Vector Store</description>
<url>https://github.com/spring-projects-experimental/spring-ai</url>
<scm>
<url>https://github.com/spring-projects-experimental/spring-ai</url>
<connection>git://github.com/spring-projects-experimental/spring-ai.git</connection>
<developerConnection>git@github.com:spring-projects-experimental/spring-ai.git</developerConnection>
</scm>
<dependencies>
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-core</artifactId>
<version>${parent.version}</version>
</dependency>
<!-- TESTING -->
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-openai</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-test</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>testcontainers</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>junit-jupiter</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>
</dependencies>
<!-- <build>
<plugins>
<plugin>
<groupId>io.swagger.codegen.v3</groupId>
<artifactId>swagger-codegen-maven-plugin</artifactId>
<version>3.0.50</version>
<executions>
<execution>
<goals>
<goal>generate</goal>
</goals>
<configuration>
<inputSpec>src/test/resources/api.yaml</inputSpec>
<language>java</language>
<configOptions>
<sourceFolder>src/gen/java/main</sourceFolder>
</configOptions>
</configuration>
</execution>
</executions>
</plugin>
</plugins>
</build> -->
</project>

380
vector-stores/spring-ai-chroma/src/main/java/org/springframework/experimental/ai/chroma/ChromaApi.java Normal file
View File

@@ -0,0 +1,380 @@
/*
* 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.experimental.ai.chroma;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.experimental.ai.chroma.ChromaApi.QueryRequest.Include;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.client.support.BasicAuthenticationInterceptor;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.client.HttpServerErrorException;
import org.springframework.web.client.RestTemplate;
/**
* Single-class Chroma API implementation based on the (unofficial) Chroma REST API.
*
* @author Christian Tzolov
*/
public class ChromaApi {
// Regular expression pattern that looks for a message inside the ValueError(...).
private static Pattern VALUE_ERROR_PATTERN = Pattern.compile("ValueError\\('([^']*)'\\)");
private final String baseUrl;
private final RestTemplate restTemplate;
private final ObjectMapper objectMapper;
private String keyToken;
public ChromaApi(String baseUrl, RestTemplate restTemplate) {
this(baseUrl, restTemplate, new ObjectMapper());
}
public ChromaApi(String baseUrl, RestTemplate restTemplate, ObjectMapper objectMapper) {
this.baseUrl = baseUrl;
this.restTemplate = restTemplate;
this.objectMapper = objectMapper;
}
/**
* Configure access to ChromaDB secured with static API Token Authentication:
* https://docs.trychroma.com/usage-guide#static-api-token-authentication
* @param keyToken Chroma static API Token Authentication. (Optional)
*/
public ChromaApi withKeyToken(String keyToken) {
this.keyToken = keyToken;
return this;
}
/**
* Configure access to ChromaDB secured with Basic Authentication:
* https://docs.trychroma.com/usage-guide#basic-authentication
* @param username Credentials username.
* @param password Credentials password.
*/
public ChromaApi withBasicAuthCredentials(String username, String password) {
this.restTemplate.getInterceptors().add(new BasicAuthenticationInterceptor(username, username));
return this;
}
/**
* Chroma embedding collection.
*
* @param id Collection Id.
* @param name The name of the collection.
* @param metadata Metadata associated with the collection.
*/
public record Collection(String id, String name, Map<String, String> metadata) {
}
/**
* Request to create a new collection with the given name and metadata.
*
* @param name The name of the collection to create.
* @param metadata Optional metadata to associate with the collection.
*/
public record CreateCollectionRequest(String name, Map<String, String> metadata) {
public CreateCollectionRequest(String name) {
this(name, new HashMap<>(Map.of("hnsw:space", "cosine")));
}
}
/**
* Add embeddings to the chroma data store.
*
* @param ids The ids of the embeddings to add.
* @param embeddings The embeddings to add.
* @param metadata The metadata to associate with the embeddings. When querying, you
* can filter on this metadata.
* @param documents The documents contents to associate with the embeddings.
*/
public record AddEmbeddingsRequest(List<String> ids, List<float[]> embeddings,
@JsonProperty("metadatas") List<Map<String, Object>> metadata, List<String> documents) {
// Convenance for adding a single embedding.
public AddEmbeddingsRequest(String id, float[] embedding, Map<String, Object> metadata, String document) {
this(List.of(id), List.of(embedding), List.of(metadata), List.of(document));
}
}
/**
* Request to delete embedding from a collection.
*
* @param ids The ids of the embeddings to delete. (Optional)
* @param where Condition to filter items to delete based on metadata values.
* (Optional)
*/
public record DeleteEmbeddingsRequest(List<String> ids, Map<String, Object> where) {
public DeleteEmbeddingsRequest(List<String> ids) {
this(ids, Map.of());
}
}
/**
* Get embeddings from a collection.
*
* @param ids IDs of the embeddings to get.
* @param where Condition to filter results based on metadata values.
* @param limit Limit on the number of collection embeddings to get.
* @param offset Offset on the embeddings to get.
* @param include A list of what to include in the results. Can contain "embeddings",
* "metadatas", "documents", "distances". Ids are always included. Defaults to
* [metadatas, documents, distances].
*/
public record GetEmbeddingsRequest(List<String> ids, Map<String, Object> where, int limit, int offset,
List<Include> include) {
public GetEmbeddingsRequest(List<String> ids) {
this(ids, Map.of(), 10, 0, Include.all);
}
public GetEmbeddingsRequest(List<String> ids, Map<String, Object> where) {
this(ids, where, 10, 0, Include.all);
}
public GetEmbeddingsRequest(List<String> ids, Map<String, Object> where, int limit, int offset) {
this(ids, where, limit, offset, Include.all);
}
}
/**
* Object containing the get embedding results.
*
* @param ids List of document ids. One for each returned document.
* @param embeddings List of document embeddings. One for each returned document.
* @param documents List of document contents. One for each returned document.
* @param metadata List of document metadata. One for each returned document.
*/
public record GetEmbeddingResponse(List<String> ids, List<List<Float>> embeddings, List<String> documents,
@JsonProperty("metadatas") List<Map<String, String>> metadata) {
}
/**
* Request to get the nResults nearest neighbor embeddings for provided
* queryEmbeddings.
*
* @param queryEmbeddings The embeddings to get the closes neighbors of.
* @param nResults The number of neighbors to return for each query_embedding or
* query_texts.
* @param where Condition to filter results based on metadata values.
* @param include A list of what to include in the results. Can contain "embeddings",
* "metadatas", "documents", "distances". Ids are always included. Defaults to
* [metadatas, documents, distances].
*/
public record QueryRequest(@JsonProperty("query_embeddings") List<List<Float>> queryEmbeddings,
@JsonProperty("n_results") int nResults, Map<String, Object> where, List<Include> include) {
public enum Include {
metadatas, documents, distances, embeddings;
public static final List<Include> all = List.of(metadatas, documents, distances, embeddings);
}
/**
* Convenience to query for a single embedding instead of a batch of embeddings.
*/
public QueryRequest(List<Float> queryEmbedding, int nResults) {
this(List.of(queryEmbedding), nResults, Map.of(), Include.all);
}
public QueryRequest(List<Float> queryEmbedding, int nResults, Map<String, Object> where) {
this(List.of(queryEmbedding), nResults, where, Include.all);
}
}
/**
* A QueryResponse object containing the query results.
*
* @param ids List of list of document ids. One for each returned document.
* @param embeddings List of list of document embeddings. One for each returned
* document.
* @param documents List of list of document contents. One for each returned document.
* @param metadata List of list of document metadata. One for each returned document.
* @param distances List of list of search distances. One for each returned document.
*/
public record QueryResponse(List<List<String>> ids, List<List<List<Float>>> embeddings,
List<List<String>> documents, @JsonProperty("metadatas") List<List<Map<String, Object>>> metadata,
List<List<Double>> distances) {
}
/**
* Single query embedding response.
*/
public record Embedding(String id, List<Float> embedding, String document, Map<String, Object> metadata,
Double distances) {
}
public List<Embedding> toEmbeddingResponseList(QueryResponse queryResponse) {
List<Embedding> result = new ArrayList<>();
if (queryResponse != null && !CollectionUtils.isEmpty(queryResponse.ids())) {
for (int i = 0; i < queryResponse.ids().get(0).size(); i++) {
result.add(new Embedding(queryResponse.ids().get(0).get(i), queryResponse.embeddings().get(0).get(i),
queryResponse.documents().get(0).get(i), queryResponse.metadata().get(0).get(i),
queryResponse.distances().get(0).get(i)));
}
}
return result;
}
//
// Chroma Client API (https://docs.trychroma.com/js_reference/Client)
//
public Collection createCollection(CreateCollectionRequest createCollectionRequest) {
return this.restTemplate
.exchange(this.baseUrl + "/api/v1/collections", HttpMethod.POST,
this.getHttpEntityFor(createCollectionRequest), Collection.class)
.getBody();
}
/**
* Delete a collection with the given name.
* @param collectionName the name of the collection to delete.
*
*/
public void deleteCollection(String collectionName) {
this.restTemplate.exchange(this.baseUrl + "/api/v1/collections/{collection_name}", HttpMethod.DELETE,
new HttpEntity<>(httpHeaders()), Void.class, collectionName);
}
public Collection getCollection(String collectionName) {
try {
return this.restTemplate
.exchange(this.baseUrl + "/api/v1/collections/{collection_name}", HttpMethod.GET,
new HttpEntity<>(httpHeaders()), Collection.class, collectionName)
.getBody();
}
catch (HttpServerErrorException e) {
String msg = this.getValueErrorMessage(e.getMessage());
if (String.format("Collection %s does not exist.", collectionName).equals(msg)) {
return null;
}
throw new RuntimeException(msg, e);
}
}
private static class CollectionList extends ArrayList<Collection> {
}
public List<Collection> listCollections() {
return this.restTemplate
.exchange(this.baseUrl + "/api/v1/collections/", HttpMethod.GET, new HttpEntity<>(httpHeaders()),
CollectionList.class)
.getBody();
}
//
// Chroma Collection API (https://docs.trychroma.com/js_reference/Collection)
//
public Boolean upsertEmbeddings(String collectionId, AddEmbeddingsRequest embedding) {
return this.restTemplate
.exchange(this.baseUrl + "/api/v1/collections/{collection_id}/upsert", HttpMethod.POST,
this.getHttpEntityFor(embedding), Boolean.class, collectionId)
.getBody();
}
public List<String> deleteEmbeddings(String collectionId, DeleteEmbeddingsRequest deleteRequest) {
return this.restTemplate
.exchange(this.baseUrl + "/api/v1/collections/{collection_id}/delete", HttpMethod.POST,
this.getHttpEntityFor(deleteRequest), List.class, collectionId)
.getBody();
}
public Long countEmbeddings(String collectionId) {
return this.restTemplate
.exchange(this.baseUrl + "/api/v1/collections/{collection_id}/count", HttpMethod.GET,
new HttpEntity<>(httpHeaders()), Long.class, collectionId)
.getBody();
}
public QueryResponse queryCollection(String collectionId, QueryRequest queryRequest) {
return this.restTemplate
.exchange(this.baseUrl + "/api/v1/collections/{collection_id}/query", HttpMethod.POST,
this.getHttpEntityFor(queryRequest), QueryResponse.class, collectionId)
.getBody();
}
public GetEmbeddingResponse getEmbeddings(String collectionId, GetEmbeddingsRequest getEmbeddingsRequest) {
return this.restTemplate
.exchange(this.baseUrl + "/api/v1/collections/{collection_id}/get", HttpMethod.POST,
this.getHttpEntityFor(getEmbeddingsRequest), GetEmbeddingResponse.class, collectionId)
.getBody();
}
// Utils
public Map<String, Object> where(String text) {
try {
return this.objectMapper.readValue(text, Map.class);
}
catch (JsonProcessingException e) {
throw new RuntimeException(e);
}
}
private <T> HttpEntity<T> getHttpEntityFor(T body) {
return new HttpEntity<>(body, httpHeaders());
}
private HttpHeaders httpHeaders() {
HttpHeaders headers = new HttpHeaders();
headers.setContentType(MediaType.APPLICATION_JSON);
if (StringUtils.hasText(this.keyToken)) {
headers.setBearerAuth(this.keyToken);
}
return headers;
}
private String getValueErrorMessage(String logString) {
if (!StringUtils.hasText(logString)) {
return "";
}
Matcher m = VALUE_ERROR_PATTERN.matcher(logString);
return (m.find()) ? m.group(1) : "";
}
}

161
vector-stores/spring-ai-chroma/src/main/java/org/springframework/experimental/ai/vectorsore/ChromaVectorStore.java Normal file
View File

@@ -0,0 +1,161 @@
/*
* 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.experimental.ai.vectorsore;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.converter.ChromaFilterExpressionConverter;
import org.springframework.ai.vectorstore.filter.converter.FilterExpressionConverter;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.experimental.ai.chroma.ChromaApi;
import org.springframework.experimental.ai.chroma.ChromaApi.AddEmbeddingsRequest;
import org.springframework.experimental.ai.chroma.ChromaApi.DeleteEmbeddingsRequest;
import org.springframework.experimental.ai.chroma.ChromaApi.Embedding;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
/**
* @author Christian Tzolov
*/
public class ChromaVectorStore implements VectorStore, InitializingBean {
public static final String DISTANCE_FIELD_NAME = "distance";
public static final String DEFAULT_COLLECTION_NAME = "SpringAiCollection";
public static final double SIMILARITY_THRESHOLD_ALL = 0.0;
public static final int DEFAULT_TOP_K = 4;
private final EmbeddingClient embeddingClient;
private final ChromaApi chromaApi;
private final String collectionName;
private FilterExpressionConverter filterExpressionConverter;
private String collectionId;
public ChromaVectorStore(EmbeddingClient embeddingClient, ChromaApi chromaApi) {
this(embeddingClient, chromaApi, DEFAULT_COLLECTION_NAME);
}
public ChromaVectorStore(EmbeddingClient embeddingClient, ChromaApi chromaApi, String collectionName) {
this.embeddingClient = embeddingClient;
this.chromaApi = chromaApi;
this.collectionName = collectionName;
this.filterExpressionConverter = new ChromaFilterExpressionConverter();
}
public void setFilterExpressionConverter(FilterExpressionConverter filterExpressionConverter) {
Assert.notNull(filterExpressionConverter, "FilterExpressionConverter should not be null.");
this.filterExpressionConverter = filterExpressionConverter;
}
@Override
public void add(List<Document> documents) {
Assert.notNull(documents, "Documents must not be null");
if (CollectionUtils.isEmpty(documents)) {
return;
}
List<String> ids = new ArrayList<>();
List<Map<String, Object>> metadatas = new ArrayList<>();
List<String> contents = new ArrayList<>();
List<float[]> embeddings = new ArrayList<>();
for (Document document : documents) {
ids.add(document.getId());
metadatas.add(document.getMetadata());
contents.add(document.getContent());
document.setEmbedding(this.embeddingClient.embed(document));
embeddings.add(JsonUtils.toFloatArray(document.getEmbedding()));
}
var success = this.chromaApi.upsertEmbeddings(this.collectionId,
new AddEmbeddingsRequest(ids, embeddings, metadatas, contents));
if (!success) {
throw new RuntimeException("Unsuccessful storing!");
}
}
@Override
public Optional<Boolean> delete(List<String> idList) {
Assert.notNull(idList, "Document id list must not be null");
List<String> deletedIds = this.chromaApi.deleteEmbeddings(this.collectionId,
new DeleteEmbeddingsRequest(idList));
return Optional.of(deletedIds.size() == idList.size());
}
@Override
public List<Document> similaritySearch(SearchRequest request) {
String nativeFilterExpression = (request.getFilterExpression() != null)
? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : "";
String query = request.getQuery();
Assert.notNull(query, "Query string must not be null");
List<Double> embedding = this.embeddingClient.embed(query);
Map<String, Object> where = (StringUtils.hasText(nativeFilterExpression))
? JsonUtils.jsonToMap(nativeFilterExpression) : Map.of();
var queryRequest = new ChromaApi.QueryRequest(JsonUtils.toFloatList(embedding), request.getTopK(), where);
var queryResponse = this.chromaApi.queryCollection(this.collectionId, queryRequest);
var embeddings = this.chromaApi.toEmbeddingResponseList(queryResponse);
List<Document> responseDocuments = new ArrayList<>();
for (Embedding chromaEmbedding : embeddings) {
float distance = chromaEmbedding.distances().floatValue();
if ((1 - distance) >= request.getSimilarityThreshold()) {
String id = chromaEmbedding.id();
String content = chromaEmbedding.document();
Map<String, Object> metadata = chromaEmbedding.metadata();
if (metadata == null) {
metadata = new HashMap<>();
}
metadata.put(DISTANCE_FIELD_NAME, distance);
Document document = new Document(id, content, metadata);
document.setEmbedding(JsonUtils.toDouble(chromaEmbedding.embedding()));
responseDocuments.add(document);
}
}
return responseDocuments;
}
@Override
public void afterPropertiesSet() throws Exception {
var collection = this.chromaApi.getCollection(this.collectionName);
if (collection == null) {
collection = this.chromaApi.createCollection(new ChromaApi.CreateCollectionRequest(this.collectionName));
}
this.collectionId = collection.id();
}
}

84
vector-stores/spring-ai-chroma/src/main/java/org/springframework/experimental/ai/vectorsore/JsonUtils.java Normal file
View File

@@ -0,0 +1,84 @@
/*
* 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.experimental.ai.vectorsore;
import java.util.List;
import java.util.Map;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
/**
* @author Christian Tzolov
*/
/**
* Utility class for JSON processing. Provides methods for converting JSON strings to maps
* and lists, and for converting between lists of different numeric types.
*
* @author Christian Tzolov
*/
public class JsonUtils {
/**
* Converts a JSON string to a map.
* @param jsonText the JSON string to convert
* @return the map representation of the JSON string
* @throws RuntimeException if an error occurs during the conversion
*/
public static Map<String, Object> jsonToMap(String jsonText) {
try {
return (Map<String, Object>) new ObjectMapper().readValue(jsonText, Map.class);
}
catch (JsonProcessingException e) {
throw new RuntimeException(e);
}
}
/**
* Converts a list of doubles to a list of floats.
* @param embeddingDouble the list of doubles to convert
* @return the list of floats
*/
public static List<Float> toFloatList(List<Double> embeddingDouble) {
return embeddingDouble.stream().map(Number::floatValue).toList();
}
/**
* Converts a list of doubles to a float array.
* @param embeddingDouble the list of doubles to convert
* @return the float array
*/
public static float[] toFloatArray(List<Double> embeddingDouble) {
float[] embeddingFloat = new float[embeddingDouble.size()];
int i = 0;
for (Double d : embeddingDouble) {
embeddingFloat[i++] = d.floatValue();
}
return embeddingFloat;
}
/**
* Converts a list of floats to a list of doubles.
* @param floats the list of floats to convert
* @return the list of doubles
*/
public static List<Double> toDouble(List<Float> floats) {
return floats.stream().map(f -> f.doubleValue()).toList();
}
}

186
vector-stores/spring-ai-chroma/src/test/java/org/springframework/experimental/ai/chroma/ChromaApiIT.java Normal file
View File

@@ -0,0 +1,186 @@
/*
* 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.experimental.ai.chroma;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
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.experimental.ai.chroma.ChromaApi.AddEmbeddingsRequest;
import org.springframework.experimental.ai.chroma.ChromaApi.Collection;
import org.springframework.experimental.ai.chroma.ChromaApi.GetEmbeddingsRequest;
import org.springframework.experimental.ai.chroma.ChromaApi.QueryRequest;
import org.springframework.web.client.RestTemplate;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
*/
@SpringBootTest
@Testcontainers
public class ChromaApiIT {
@Container
static GenericContainer<?> chromaContainer = new GenericContainer<>("ghcr.io/chroma-core/chroma:0.4.15")
.withExposedPorts(8000);
@Autowired
ChromaApi chroma;
@BeforeEach
public void beforeEach() {
chroma.listCollections().stream().forEach(c -> chroma.deleteCollection(c.name()));
}
@Test
public void testClient() {
var newCollection = chroma.createCollection(new ChromaApi.CreateCollectionRequest("TestCollection"));
assertThat(newCollection).isNotNull();
assertThat(newCollection.name()).isEqualTo("TestCollection");
var getCollection = chroma.getCollection("TestCollection");
assertThat(getCollection).isNotNull();
assertThat(getCollection.name()).isEqualTo("TestCollection");
assertThat(getCollection.id()).isEqualTo(newCollection.id());
List<Collection> collections = chroma.listCollections();
assertThat(collections).hasSize(1);
assertThat(collections.get(0).id()).isEqualTo(newCollection.id());
chroma.deleteCollection(newCollection.name());
assertThat(chroma.listCollections()).hasSize(0);
}
@Test
public void testCollection() {
var newCollection = chroma.createCollection(new ChromaApi.CreateCollectionRequest("TestCollection"));
assertThat(chroma.countEmbeddings(newCollection.id())).isEqualTo(0);
var addEmbeddingRequest = new AddEmbeddingsRequest(List.of("id1", "id2"),
List.of(new float[] { 1f, 1f, 1f }, new float[] { 2f, 2f, 2f }),
List.of(Map.of(), Map.of("key1", "value1", "key2", true, "key3", 23.4)),
List.of("Hello World", "Big World"));
var success = chroma.upsertEmbeddings(newCollection.id(), addEmbeddingRequest);
assertThat(success).isTrue();
var addEmbeddingRequest2 = new AddEmbeddingsRequest("id3", new float[] { 3f, 3f, 3f },
Map.of("key1", "value1", "key2", true, "key3", 23.4), "Big World");
chroma.upsertEmbeddings(newCollection.id(), addEmbeddingRequest2);
assertThat(chroma.countEmbeddings(newCollection.id())).isEqualTo(3);
// Update existing embedding.
chroma.upsertEmbeddings(newCollection.id(), new AddEmbeddingsRequest("id3", new float[] { 6f, 6f, 6f },
Map.of("key1", "value2", "key2", false, "key4", 23.4), "Small World"));
var result = chroma.getEmbeddings(newCollection.id(), new GetEmbeddingsRequest(List.of("id2")));
assertThat(result.ids().get(0)).isEqualTo("id2");
result = chroma.getEmbeddings(newCollection.id(), new GetEmbeddingsRequest(List.of(), chroma.where("""
{ "key2" : { "$eq": true} }
""")));
assertThat(result.ids()).containsExactlyInAnyOrder("id2");
}
@Test
public void testQueryWhere() {
var collection = chroma.createCollection(new ChromaApi.CreateCollectionRequest("TestCollection"));
var add1 = new AddEmbeddingsRequest("id1", new float[] { 1f, 1f, 1f },
Map.of("country", "BG", "active", true, "price", 23.4, "year", 2020),
"The World is Big and Salvation Lurks Around the Corner");
var add2 = new AddEmbeddingsRequest("id2", new float[] { 1f, 1f, 1f }, Map.of("country", "NL"),
"The World is Big and Salvation Lurks Around the Corner");
var add3 = new AddEmbeddingsRequest("id3", new float[] { 1f, 1f, 1f },
Map.of("country", "BG", "active", false, "price", 40.1, "year", 2023),
"The World is Big and Salvation Lurks Around the Corner");
chroma.upsertEmbeddings(collection.id(), add1);
chroma.upsertEmbeddings(collection.id(), add2);
chroma.upsertEmbeddings(collection.id(), add3);
assertThat(chroma.countEmbeddings(collection.id())).isEqualTo(3);
var queryResult = chroma.queryCollection(collection.id(), new QueryRequest(List.of(1f, 1f, 1f), 3));
assertThat(queryResult.ids().get(0)).hasSize(3);
assertThat(queryResult.ids().get(0)).containsExactlyInAnyOrder("id1", "id2", "id3");
var chromaEmbeddings = chroma.toEmbeddingResponseList(queryResult);
assertThat(chromaEmbeddings).hasSize(3);
assertThat(chromaEmbeddings).hasSize(3);
queryResult = chroma.queryCollection(collection.id(), new QueryRequest(List.of(1f, 1f, 1f), 3, chroma.where("""
{
"$and" : [
{"country" : { "$eq": "BG"}},
{"year" : { "$gte": 2020}}
]
}
""")));
assertThat(queryResult.ids().get(0)).hasSize(2);
assertThat(queryResult.ids().get(0)).containsExactlyInAnyOrder("id1", "id3");
queryResult = chroma.queryCollection(collection.id(), new QueryRequest(List.of(1f, 1f, 1f), 3, chroma.where("""
{
"$and" : [
{"country" : { "$eq": "BG"}},
{"year" : { "$gte": 2020}},
{"active" : { "$eq": true}}
]
}
""")));
assertThat(queryResult.ids().get(0)).hasSize(1);
assertThat(queryResult.ids().get(0)).containsExactlyInAnyOrder("id1");
}
@SpringBootConfiguration
public static class Config {
@Bean
public RestTemplate restTemplate() {
return new RestTemplate();
}
@Bean
public ChromaApi chromaApi(RestTemplate restTemplate) {
int port = chromaContainer.getMappedPort(8000);
return new ChromaApi("http://localhost:" + port, restTemplate);
}
}
}

141
vector-stores/spring-ai-chroma/src/test/java/org/springframework/experimental/ai/vectorstore/BasicAuthChromaWhereIT.java Normal file
View File

@@ -0,0 +1,141 @@
/*
* 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.experimental.ai.vectorstore;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import com.theokanning.openai.client.OpenAiApi;
import com.theokanning.openai.service.OpenAiService;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.images.builder.Transferable;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.experimental.ai.chroma.ChromaApi;
import org.springframework.experimental.ai.vectorsore.ChromaVectorStore;
import org.springframework.web.client.RestTemplate;
import static org.assertj.core.api.Assertions.assertThat;
/**
* ChromaDB with Basic Authentication:
* https://docs.trychroma.com/usage-guide#basic-authentication
*
* The scr/test/resource/server.htpasswd file is generated with:
* <code>htpasswd -Bbn admin admin > server.htpasswd</code>
*
* @author Christian Tzolov
*/
@Testcontainers
public class BasicAuthChromaWhereIT {
public static String CHROMA_SERVER_URL = "http://localhost:";
/**
* ChromaDB with Basic Authentication:
* https://docs.trychroma.com/usage-guide#basic-authentication
*/
@Container
static GenericContainer<?> chromaContainer = new GenericContainer<>("ghcr.io/chroma-core/chroma:0.4.15")
.withEnv("CHROMA_SERVER_AUTH_CREDENTIALS_FILE", "server.htpasswd")
.withEnv("CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER",
"chromadb.auth.providers.HtpasswdFileServerAuthCredentialsProvider")
.withEnv("CHROMA_SERVER_AUTH_PROVIDER", "chromadb.auth.basic.BasicAuthServerProvider")
.withCopyToContainer(Transferable.of("src/test/resources/server.htpasswd"), "server.htpasswd")
.withExposedPorts(8000);
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class)
.withPropertyValues("spring.ai.openai.apiKey=" + System.getenv("OPENAI_API_KEY"));
@Test
public void withInFiltersExpressions1() {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(List.of(new Document("1", "Article by john", Map.of("author", "john")),
new Document("2", "Article by Jack", Map.of("author", "jack")),
new Document("3", "Article by Jill", Map.of("author", "jill"))));
String query = "Give me articles by john";
List<Document> results = vectorStore.similaritySearch(SearchRequest.query(query).withTopK(5));
assertThat(results).hasSize(3);
results = vectorStore.similaritySearch(SearchRequest.query(query)
.withTopK(5)
.withSimilarityThresholdAll()
.withFilterExpression("author in ['john', 'jill']"));
assertThat(results).hasSize(2);
assertThat(results.stream().map(d -> d.getId()).toList()).containsExactlyInAnyOrder("1", "3");
});
}
@SpringBootConfiguration
public static class TestApplication {
@Bean
public RestTemplate restTemplate() {
return new RestTemplate();
}
@Bean
public ChromaApi chromaApi(RestTemplate restTemplate) {
int port = chromaContainer.getMappedPort(8000);
return new ChromaApi(CHROMA_SERVER_URL + port, restTemplate).withBasicAuthCredentials("admin", "admin");
}
@Bean
public VectorStore chromaVectorStore(EmbeddingClient embeddingClient, ChromaApi chromaApi) {
return new ChromaVectorStore(embeddingClient, chromaApi, "TestCollection");
}
@Bean
public EmbeddingClient embeddingClient() {
Retrofit retrofit = new Retrofit.Builder().baseUrl("https://api.openai.com")
.client(OpenAiService.defaultClient(System.getenv("OPENAI_API_KEY"), Duration.ofSeconds(60)))
.addConverterFactory(JacksonConverterFactory.create(OpenAiService.defaultObjectMapper()))
.addCallAdapterFactory(RxJava2CallAdapterFactory.create())
.build();
OpenAiApi api = retrofit.create(OpenAiApi.class);
return new OpenAiEmbeddingClient(new OpenAiService(api), "text-embedding-ada-002");
}
}
}

240
vector-stores/spring-ai-chroma/src/test/java/org/springframework/experimental/ai/vectorstore/ChromaVectorStoreIT.java Normal file
View File

@@ -0,0 +1,240 @@
/*
* 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.experimental.ai.vectorstore;
import java.time.Duration;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import com.theokanning.openai.client.OpenAiApi;
import com.theokanning.openai.service.OpenAiService;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.experimental.ai.chroma.ChromaApi;
import org.springframework.experimental.ai.vectorsore.ChromaVectorStore;
import org.springframework.web.client.RestTemplate;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
*/
@Testcontainers
public class ChromaVectorStoreIT {
@Container
static GenericContainer<?> chromaContainer = new GenericContainer<>("ghcr.io/chroma-core/chroma:0.4.15")
.withExposedPorts(8000);
List<Document> documents = List.of(
new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!",
Collections.singletonMap("meta1", "meta1")),
new Document("Hello World Hello World Hello World Hello World Hello World Hello World Hello World"),
new Document(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression",
Collections.singletonMap("meta2", "meta2")));
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class)
.withPropertyValues("spring.ai.openai.apiKey=" + System.getenv("OPENAI_API_KEY"));
@Test
public void addAndSearch() {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Great").withTopK(1));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getContent()).isEqualTo(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression");
assertThat(resultDoc.getMetadata()).containsKeys("meta2", "distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
List<Document> results2 = vectorStore.similaritySearch(SearchRequest.query("Great").withTopK(1));
assertThat(results2).hasSize(0);
});
}
@Test
public void addAndSearchWithFilters() {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "Bulgaria"));
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "Netherland"));
vectorStore.add(List.of(bgDocument, nlDocument));
var request = SearchRequest.query("The World").withTopK(5);
List<Document> results = vectorStore.similaritySearch(request);
assertThat(results).hasSize(2);
results = vectorStore
.similaritySearch(request.withSimilarityThresholdAll().withFilterExpression("country == 'Bulgaria'"));
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId());
results = vectorStore
.similaritySearch(request.withSimilarityThresholdAll().withFilterExpression("country == 'Netherland'"));
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId());
// Remove all documents from the store
vectorStore.delete(List.of(bgDocument, nlDocument).stream().map(doc -> doc.getId()).toList());
});
}
@Test
public void documentUpdateTest() {
// Note ,using OpenAI to calculate embeddings
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
Collections.singletonMap("meta1", "meta1"));
vectorStore.add(List.of(document));
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(5));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getContent()).isEqualTo("Spring AI rocks!!");
assertThat(resultDoc.getMetadata()).containsKey("meta1");
assertThat(resultDoc.getMetadata()).containsKey("distance");
Document sameIdDocument = new Document(document.getId(),
"The World is Big and Salvation Lurks Around the Corner",
Collections.singletonMap("meta2", "meta2"));
vectorStore.add(List.of(sameIdDocument));
results = vectorStore.similaritySearch(SearchRequest.query("FooBar").withTopK(5));
assertThat(results).hasSize(1);
resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getContent()).isEqualTo("The World is Big and Salvation Lurks Around the Corner");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
// Remove all documents from the store
vectorStore.delete(List.of(document.getId()));
});
}
@Test
public void searchThresholdTest() {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
var request = SearchRequest.query("Great").withTopK(5);
List<Document> fullResult = vectorStore.similaritySearch(request.withSimilarityThresholdAll());
List<Float> distances = fullResult.stream().map(doc -> (Float) doc.getMetadata().get("distance")).toList();
assertThat(distances).hasSize(3);
float threshold = (distances.get(0) + distances.get(1)) / 2;
List<Document> results = vectorStore.similaritySearch(request.withSimilarityThreshold(1 - threshold));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
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");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
});
}
@SpringBootConfiguration
public static class TestApplication {
@Bean
public RestTemplate restTemplate() {
return new RestTemplate();
}
@Bean
public ChromaApi chromaApi(RestTemplate restTemplate) {
int port = chromaContainer.getMappedPort(8000);
return new ChromaApi("http://localhost:" + port, restTemplate);
}
@Bean
public VectorStore chromaVectorStore(EmbeddingClient embeddingClient, ChromaApi chromaApi) {
return new ChromaVectorStore(embeddingClient, chromaApi, "TestCollection");
}
@Bean
public EmbeddingClient embeddingClient() {
Retrofit retrofit = new Retrofit.Builder().baseUrl("https://api.openai.com")
.client(OpenAiService.defaultClient(System.getenv("OPENAI_API_KEY"), Duration.ofSeconds(60)))
.addConverterFactory(JacksonConverterFactory.create(OpenAiService.defaultObjectMapper()))
.addCallAdapterFactory(RxJava2CallAdapterFactory.create())
.build();
OpenAiApi api = retrofit.create(OpenAiApi.class);
return new OpenAiEmbeddingClient(new OpenAiService(api), "text-embedding-ada-002");
}
}
}

174
vector-stores/spring-ai-chroma/src/test/java/org/springframework/experimental/ai/vectorstore/TokenSecuredChromaWhereIT.java Normal file
View File

@@ -0,0 +1,174 @@
/*
* 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.experimental.ai.vectorstore;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import com.theokanning.openai.client.OpenAiApi;
import com.theokanning.openai.service.OpenAiService;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.experimental.ai.chroma.ChromaApi;
import org.springframework.experimental.ai.vectorsore.ChromaVectorStore;
import org.springframework.web.client.RestTemplate;
import static org.assertj.core.api.Assertions.assertThat;
/**
* ChromaDB with static API Token Authentication:
* https://docs.trychroma.com/usage-guide#static-api-token-authentication
*
* Test cases are based on the Chroma:
* https://docs.trychroma.com/usage-guide#using-where-filters and the related
* https://github.com/chroma-core/chroma/blob/main/examples/basic_functionality/in_not_in_filtering.ipynb
*
* @author Christian Tzolov
*/
@Testcontainers
public class TokenSecuredChromaWhereIT {
public static String CHROMA_SERVER_URL = "http://localhost:";
public static String CHROMA_SERVER_AUTH_CREDENTIALS = "test-token";
/**
* ChromaDB with static API Token Authentication:
* https://docs.trychroma.com/usage-guide#static-api-token-authentication
*/
@Container
static GenericContainer<?> chromaContainer = new GenericContainer<>("ghcr.io/chroma-core/chroma:0.4.15")
.withEnv("CHROMA_SERVER_AUTH_CREDENTIALS", CHROMA_SERVER_AUTH_CREDENTIALS)
.withEnv("CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER",
"chromadb.auth.token.TokenConfigServerAuthCredentialsProvider")
.withEnv("CHROMA_SERVER_AUTH_PROVIDER", "chromadb.auth.token.TokenAuthServerProvider")
.withExposedPorts(8000);
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class)
.withPropertyValues("spring.ai.openai.apiKey=" + System.getenv("OPENAI_API_KEY"));
@Test
public void withInFiltersExpressions1() {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(List.of(new Document("1", "Article by john", Map.of("author", "john")),
new Document("2", "Article by Jack", Map.of("author", "jack")),
new Document("3", "Article by Jill", Map.of("author", "jill"))));
var request = SearchRequest.query("Give me articles by john").withTopK(5);
List<Document> results = vectorStore.similaritySearch(request);
assertThat(results).hasSize(3);
results = vectorStore.similaritySearch(
request.withSimilarityThresholdAll().withFilterExpression("author in ['john', 'jill']"));
assertThat(results).hasSize(2);
assertThat(results.stream().map(d -> d.getId()).toList()).containsExactlyInAnyOrder("1", "3");
});
}
@Test
public void withInFiltersExpressions() {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore
.add(List.of(new Document("1", "Article by john", Map.of("author", "john", "article_type", "blog")),
new Document("2", "Article by Jack", Map.of("author", "jack", "article_type", "social")),
new Document("3", "Article by Jill", Map.of("author", "jill", "article_type", "paper"))));
var request = SearchRequest.query("Give me articles by john").withTopK(5);
List<Document> results = vectorStore.similaritySearch(request);
assertThat(results).hasSize(3);
results = vectorStore.similaritySearch(request.withSimilarityThresholdAll()
.withFilterExpression("author in ['john', 'jill'] && 'article_type' == 'blog'"));
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo("1");
results = vectorStore.similaritySearch(request.withSimilarityThresholdAll()
.withFilterExpression("author in ['john'] || 'article_type' == 'paper'"));
assertThat(results).hasSize(2);
assertThat(results.stream().map(d -> d.getId()).toList()).containsExactlyInAnyOrder("1", "3");
});
}
@SpringBootConfiguration
public static class TestApplication {
@Bean
public RestTemplate restTemplate() {
return new RestTemplate();
}
@Bean
public ChromaApi chromaApi(RestTemplate restTemplate) {
int port = chromaContainer.getMappedPort(8000);
var chromaApi = new ChromaApi(CHROMA_SERVER_URL + port, restTemplate);
chromaApi.withKeyToken(CHROMA_SERVER_AUTH_CREDENTIALS);
return chromaApi;
}
@Bean
public VectorStore chromaVectorStore(EmbeddingClient embeddingClient, ChromaApi chromaApi) {
return new ChromaVectorStore(embeddingClient, chromaApi, "TestCollection");
}
@Bean
public EmbeddingClient embeddingClient() {
Retrofit retrofit = new Retrofit.Builder().baseUrl("https://api.openai.com")
.client(OpenAiService.defaultClient(System.getenv("OPENAI_API_KEY"), Duration.ofSeconds(60)))
.addConverterFactory(JacksonConverterFactory.create(OpenAiService.defaultObjectMapper()))
.addCallAdapterFactory(RxJava2CallAdapterFactory.create())
.build();
OpenAiApi api = retrofit.create(OpenAiApi.class);
return new OpenAiEmbeddingClient(new OpenAiService(api), "text-embedding-ada-002");
}
}
}

630
vector-stores/spring-ai-chroma/src/test/resources/api.yaml Normal file
View File

@@ -0,0 +1,630 @@
openapi: "3.0.0"
info:
title: FastAPI
version: 0.1.0
paths:
/api/v1:
get:
summary: Root
operationId: root
responses:
'200':
description: Successful Response
content:
application/json:
schema:
additionalProperties:
type: integer
type: object
title: Response Root Api V1 Get
/api/v1/reset:
post:
summary: Reset
operationId: reset
responses:
'200':
description: Successful Response
content:
application/json:
schema:
type: boolean
title: Response Reset Api V1 Reset Post
/api/v1/version:
get:
summary: Version
operationId: version
responses:
'200':
description: Successful Response
content:
application/json:
schema:
type: string
title: Response Version Api V1 Version Get
/api/v1/heartbeat:
get:
summary: Heartbeat
operationId: heartbeat
responses:
'200':
description: Successful Response
content:
application/json:
schema:
additionalProperties:
type: number
type: object
title: Response Heartbeat Api V1 Heartbeat Get
/api/v1/raw_sql:
post:
summary: Raw Sql
operationId: raw_sql
requestBody:
content:
application/json:
schema:
$ref: '#/components/schemas/RawSql'
required: true
responses:
'200':
description: Successful Response
content:
application/json:
schema: {}
'422':
description: Validation Error
content:
application/json:
schema:
$ref: '#/components/schemas/HTTPValidationError'
/api/v1/collections:
get:
summary: List Collections
operationId: list_collections
responses:
'200':
description: Successful Response
content:
application/json:
schema: {}
post:
summary: Create Collection
operationId: create_collection
requestBody:
content:
application/json:
schema:
$ref: '#/components/schemas/CreateCollection'
required: true
responses:
'200':
description: Successful Response
content:
application/json:
schema: {}
'422':
description: Validation Error
content:
application/json:
schema:
$ref: '#/components/schemas/HTTPValidationError'
/api/v1/collections/{collection_id}/add:
post:
summary: Add
operationId: add
parameters:
- required: true
schema:
type: string
title: Collection Id
name: collection_id
in: path
requestBody:
content:
application/json:
schema:
$ref: '#/components/schemas/AddEmbedding'
required: true
responses:
'201':
description: Successful Response
content:
application/json:
schema: {}
'422':
description: Validation Error
content:
application/json:
schema:
$ref: '#/components/schemas/HTTPValidationError'
/api/v1/collections/{collection_id}/update:
post:
summary: Update
operationId: update
parameters:
- required: true
schema:
type: string
title: Collection Id
name: collection_id
in: path
requestBody:
content:
application/json:
schema:
$ref: '#/components/schemas/UpdateEmbedding'
required: true
responses:
'200':
description: Successful Response
content:
application/json:
schema: {}
'422':
description: Validation Error
content:
application/json:
schema:
$ref: '#/components/schemas/HTTPValidationError'
/api/v1/collections/{collection_id}/upsert:
post:
summary: Upsert
operationId: upsert
parameters:
- required: true
schema:
type: string
title: Collection Id
name: collection_id
in: path
requestBody:
content:
application/json:
schema:
$ref: '#/components/schemas/AddEmbedding'
required: true
responses:
'200':
description: Successful Response
content:
application/json:
schema: {}
'422':
description: Validation Error
content:
application/json:
schema:
$ref: '#/components/schemas/HTTPValidationError'
/api/v1/collections/{collection_id}/get:
post:
summary: Get
operationId: get
parameters:
- required: true
schema:
type: string
title: Collection Id
name: collection_id
in: path
requestBody:
content:
application/json:
schema:
$ref: '#/components/schemas/GetEmbedding'
required: true
responses:
'200':
description: Successful Response
content:
application/json:
schema: {} #TODO add actual GetResult Body
'422':
description: Validation Error
content:
application/json:
schema:
$ref: '#/components/schemas/HTTPValidationError'
/api/v1/collections/{collection_id}/delete:
post:
summary: Delete
operationId: delete
parameters:
- required: true
schema:
type: string
title: Collection Id
name: collection_id
in: path
requestBody:
content:
application/json:
schema:
$ref: '#/components/schemas/DeleteEmbedding'
required: true
responses:
'200':
description: Successful Response
content:
application/json:
schema: {}
'422':
description: Validation Error
content:
application/json:
schema:
$ref: '#/components/schemas/HTTPValidationError'
/api/v1/collections/{collection_id}/count:
get:
summary: Count
operationId: count
parameters:
- required: true
schema:
type: string
title: Collection Id
name: collection_id
in: path
responses:
'200':
description: Successful Response
content:
application/json:
schema:
type: integer
'422':
description: Validation Error
content:
application/json:
schema:
$ref: '#/components/schemas/HTTPValidationError'
/api/v1/collections/{collection_id}/query:
post:
summary: Get Nearest Neighbors
operationId: get_nearest_neighbors
parameters:
- required: true
schema:
type: string
title: Collection Id
name: collection_id
in: path
requestBody:
content:
application/json:
schema:
$ref: '#/components/schemas/QueryEmbedding'
required: true
responses:
'200':
description: Successful Response
content:
application/json:
schema: {}
'422':
description: Validation Error
content:
application/json:
schema:
$ref: '#/components/schemas/HTTPValidationError'
/api/v1/collections/{collection_name}/create_index:
post:
summary: Create Index
operationId: create_index
parameters:
- required: true
schema:
type: string
title: Collection Name
name: collection_name
in: path
responses:
'200':
description: Successful Response
content:
application/json:
schema:
type: boolean
'422':
description: Validation Error
content:
application/json:
schema:
$ref: '#/components/schemas/HTTPValidationError'
/api/v1/collections/{collection_name}:
get:
summary: Get Collection
operationId: get_collection
parameters:
- required: true
schema:
type: string
title: Collection Name
name: collection_name
in: path
responses:
'200':
description: Successful Response
content:
application/json:
schema: {}
'422':
description: Validation Error
content:
application/json:
schema:
$ref: '#/components/schemas/HTTPValidationError'
delete:
summary: Delete Collection
operationId: delete_collection
parameters:
- required: true
schema:
type: string
title: Collection Name
name: collection_name
in: path
responses:
'200':
description: Successful Response
content:
application/json:
schema: {}
'422':
description: Validation Error
content:
application/json:
schema:
$ref: '#/components/schemas/HTTPValidationError'
/api/v1/collections/{collection_id}:
put:
summary: Update Collection
operationId: update_collection
parameters:
- required: true
schema:
type: string
title: Collection Id
name: collection_id
in: path
requestBody:
content:
application/json:
schema:
$ref: '#/components/schemas/UpdateCollection'
required: true
responses:
'200':
description: Successful Response
content:
application/json:
schema: {}
'422':
description: Validation Error
content:
application/json:
schema:
$ref: '#/components/schemas/HTTPValidationError'
components:
schemas:
AddEmbedding:
properties:
embeddings:
items: {}
type: array
title: Embeddings
metadatas:
items:
type: object
additionalProperties: true
type: array
title: Metadatas
documents:
items:
type: string
type: array
title: Documents
ids:
items:
type: string
type: array
title: Ids
increment_index:
type: boolean
title: Increment Index
default: true
type: object
required:
- ids
title: AddEmbedding
CreateCollection:
properties:
name:
type: string
title: Name
metadata:
type: object
title: Metadata
get_or_create:
type: boolean
title: Get Or Create
default: false
type: object
required:
- name
title: CreateCollection
DeleteEmbedding:
properties:
ids:
items:
type: string
type: array
title: Ids
where:
type: object
title: Where
where_document:
type: object
title: Where Document
type: object
title: DeleteEmbedding
GetEmbedding:
properties:
ids:
items:
type: string
type: array
title: Ids
where:
type: object
title: Where
where_document:
type: object
title: Where Document
sort:
type: string
title: Sort
limit:
type: integer
title: Limit
offset:
type: integer
title: Offset
include:
items:
anyOf:
- type: string
enum:
- documents
- type: string
enum:
- embeddings
- type: string
enum:
- metadatas
- type: string
enum:
- distances
type: array
title: Include
default:
- metadatas
- documents
type: object
title: GetEmbedding
HTTPValidationError:
properties:
detail:
items:
$ref: '#/components/schemas/ValidationError'
type: array
title: Detail
type: object
title: HTTPValidationError
QueryEmbedding:
properties:
where:
type: object
title: Where
additionalProperties: true
default: {}
where_document:
type: object
title: Where Document
additionalProperties: true
default: {}
query_embeddings:
items: {}
type: array
additionalProperties: true
title: Query Embeddings
n_results:
type: integer
title: N Results
default: 10
include:
items:
type: string
enum:
- documents
- embeddings
- metadatas
- distances
type: array
title: Include
default:
- metadatas
- documents
- distances
type: object
required:
- query_embeddings
title: QueryEmbedding
RawSql:
properties:
raw_sql:
type: string
title: Raw Sql
type: object
required:
- raw_sql
title: RawSql
UpdateCollection:
properties:
new_name:
type: string
title: New Name
new_metadata:
type: object
title: New Metadata
type: object
title: UpdateCollection
UpdateEmbedding:
properties:
embeddings:
items: {}
type: array
title: Embeddings
metadatas:
items:
type: object
type: array
title: Metadatas
documents:
items:
type: string
type: array
title: Documents
ids:
items:
type: string
type: array
title: Ids
increment_index:
type: boolean
title: Increment Index
default: true
type: object
required:
- ids
title: UpdateEmbedding
ValidationError:
properties:
loc:
items:
anyOf:
- type: string
- type: integer
type: array
title: Location
msg:
type: string
title: Message
type:
type: string
title: Error Type
type: object
required:
- loc
- msg
- type
title: ValidationError

2
vector-stores/spring-ai-chroma/src/test/resources/server.htpasswd Normal file
View File

@@ -0,0 +1,2 @@
admin:$2y$05$fM.6b629s3L6L8RcA.kqg.BmxEzwB9t4MpGux62MEXNMJ9M7w8CY2

5
vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java
View File

@@ -54,6 +54,7 @@ import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.vectorstore.filter.converter.FilterExpressionConverter;
import org.springframework.ai.vectorstore.filter.converter.MilvusFilterExpressionConverter;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.Assert;
@@ -88,7 +89,7 @@ public class MilvusVectorStore implements VectorStore, InitializingBean {
public static final List<String> SEARCH_OUTPUT_FIELDS = Arrays.asList(DOC_ID_FIELD_NAME, CONTENT_FIELD_NAME,
METADATA_FIELD_NAME);
public final MilvusFilterExpressionConverter filterExpressionConverter = new MilvusFilterExpressionConverter();
public final FilterExpressionConverter filterExpressionConverter = new MilvusFilterExpressionConverter();
private final MilvusServiceClient milvusClient;
@@ -324,7 +325,7 @@ public class MilvusVectorStore implements VectorStore, InitializingBean {
public List<Document> similaritySearch(SearchRequest request) {
String nativeFilterExpressions = (request.getFilterExpression() != null)
? this.filterExpressionConverter.convert(request.getFilterExpression()) : "";
? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : "";
Assert.notNull(request.getQuery(), "Query string must not be null");

5
vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java
View File

@@ -33,6 +33,7 @@ import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.vectorstore.filter.converter.FilterExpressionConverter;
import org.springframework.ai.vectorstore.filter.converter.PgVectorFilterExpressionConverter;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.jdbc.core.JdbcTemplate;
@@ -56,7 +57,7 @@ public class PgVectorStore implements VectorStore, InitializingBean {
public static final String VECTOR_TABLE_NAME = "vector_store";
public final PgVectorFilterExpressionConverter filterExpressionConverter = new PgVectorFilterExpressionConverter();
public final FilterExpressionConverter filterExpressionConverter = new PgVectorFilterExpressionConverter();
private final JdbcTemplate jdbcTemplate;
@@ -273,7 +274,7 @@ public class PgVectorStore implements VectorStore, InitializingBean {
public List<Document> similaritySearch(SearchRequest request) {
String nativeFilterExpression = (request.getFilterExpression() != null)
? this.filterExpressionConverter.convert(request.getFilterExpression()) : "";
? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : "";
String jsonPathFilter = "";

12
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java
View File

@@ -34,7 +34,6 @@ import org.junit.Assert;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
@@ -42,7 +41,6 @@ import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;
import org.springframework.ai.ResourceUtils;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
@@ -58,6 +56,7 @@ import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.ApplicationContext;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Primary;
import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.util.CollectionUtils;
@@ -145,7 +144,7 @@ public class PgVectorStoreIT {
VectorStore vectorStore = context.getBean(VectorStore.class);
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "BG", "year", 2020));
Map.of("country", "BG", "year", 2020, "foo bar 1", "bar.foo"));
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "NL"));
var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner",
@@ -183,6 +182,13 @@ public class PgVectorStoreIT {
assertThat(results.get(0).getId()).isIn(bgDocument.getId(), nlDocument.getId());
assertThat(results.get(1).getId()).isIn(bgDocument.getId(), nlDocument.getId());
results = vectorStore.similaritySearch(SearchRequest.query("The World")
.withTopK(5)
.withSimilarityThresholdAll()
.withFilterExpression("\"foo bar 1\" == 'bar.foo'"));
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId());
try {
vectorStore.similaritySearch(searchRequest.withFilterExpression("country == NL"));
Assert.fail("Invalid filter expression should have been cached!");

6
vector-stores/spring-ai-pinecone/src/main/java/org/springframework/ai/vectorstore/PineconeVectorStore.java
View File

@@ -36,6 +36,8 @@ import io.pinecone.proto.Vector;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.converter.FilterExpressionConverter;
import org.springframework.ai.vectorstore.filter.converter.PineconeFilterExpressionConverter;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
@@ -53,7 +55,7 @@ public class PineconeVectorStore implements VectorStore {
private static final String DISTANCE_METADATA_FIELD_NAME = "distance";
public final PineconeFilterExpressionConverter filterExpressionConverter = new PineconeFilterExpressionConverter();
public final FilterExpressionConverter filterExpressionConverter = new PineconeFilterExpressionConverter();
private final EmbeddingClient embeddingClient;
@@ -307,7 +309,7 @@ public class PineconeVectorStore implements VectorStore {
// Filter.Expression filterExpression) {
String nativeExpressionFilters = (request.getFilterExpression() != null)
? this.filterExpressionConverter.convert(request.getFilterExpression()) : "";
? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : "";
List<Double> queryEmbedding = this.embeddingClient.embed(request.getQuery());