初始化
This commit is contained in:
56
easy-agents-store/easy-agents-store-opensearch/pom.xml
Normal file
56
easy-agents-store/easy-agents-store-opensearch/pom.xml
Normal file
@@ -0,0 +1,56 @@
|
||||
<?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/xsd/maven-4.0.0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
<parent>
|
||||
<groupId>com.easyagents</groupId>
|
||||
<artifactId>easy-agents-store</artifactId>
|
||||
<version>${revision}</version>
|
||||
</parent>
|
||||
|
||||
<name>easy-agents-store-opensearch</name>
|
||||
<artifactId>easy-agents-store-opensearch</artifactId>
|
||||
|
||||
<properties>
|
||||
<opensearch.version>2.13.0</opensearch.version>
|
||||
<httpclient5.version>5.1.4</httpclient5.version>
|
||||
<httpcore5-h2>5.1.4</httpcore5-h2>
|
||||
|
||||
<maven.compiler.source>8</maven.compiler.source>
|
||||
<maven.compiler.target>8</maven.compiler.target>
|
||||
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
||||
</properties>
|
||||
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>com.easyagents</groupId>
|
||||
<artifactId>easy-agents-core</artifactId>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.opensearch.client</groupId>
|
||||
<artifactId>opensearch-java</artifactId>
|
||||
<version>${opensearch.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.apache.httpcomponents.client5</groupId>
|
||||
<artifactId>httpclient5</artifactId>
|
||||
<version>${httpclient5.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>junit</groupId>
|
||||
<artifactId>junit</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.apache.httpcomponents.core5</groupId>
|
||||
<artifactId>httpcore5-h2</artifactId>
|
||||
<version>${httpcore5-h2}</version> <!-- 版本需与 httpclient5 一致 -->
|
||||
</dependency>
|
||||
</dependencies>
|
||||
|
||||
</project>
|
||||
@@ -0,0 +1,264 @@
|
||||
/*
|
||||
* Copyright (c) 2023-2026, Easy-Agents (fuhai999@gmail.com).
|
||||
* <p>
|
||||
* 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
|
||||
* <p>
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* <p>
|
||||
* 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 com.easyagents.store.opensearch;
|
||||
|
||||
import com.easyagents.core.document.Document;
|
||||
import com.easyagents.core.store.DocumentStore;
|
||||
import com.easyagents.core.store.SearchWrapper;
|
||||
import com.easyagents.core.store.StoreOptions;
|
||||
import com.easyagents.core.store.StoreResult;
|
||||
import com.easyagents.core.store.exception.StoreException;
|
||||
import com.easyagents.core.util.StringUtil;
|
||||
import org.apache.hc.client5.http.auth.AuthScope;
|
||||
import org.apache.hc.client5.http.auth.UsernamePasswordCredentials;
|
||||
import org.apache.hc.client5.http.impl.auth.BasicCredentialsProvider;
|
||||
import org.apache.hc.client5.http.impl.nio.PoolingAsyncClientConnectionManagerBuilder;
|
||||
import org.apache.hc.client5.http.ssl.ClientTlsStrategyBuilder;
|
||||
import org.apache.hc.client5.http.ssl.NoopHostnameVerifier;
|
||||
import org.apache.hc.core5.http.HttpHost;
|
||||
import org.apache.hc.core5.http.message.BasicHeader;
|
||||
import org.apache.hc.core5.http.nio.ssl.TlsStrategy;
|
||||
import org.apache.hc.core5.ssl.SSLContextBuilder;
|
||||
import org.opensearch.client.json.JsonData;
|
||||
import org.opensearch.client.json.jackson.JacksonJsonpMapper;
|
||||
import org.opensearch.client.opensearch.OpenSearchClient;
|
||||
import org.opensearch.client.opensearch._types.ErrorCause;
|
||||
import org.opensearch.client.opensearch._types.InlineScript;
|
||||
import org.opensearch.client.opensearch._types.mapping.Property;
|
||||
import org.opensearch.client.opensearch._types.mapping.TextProperty;
|
||||
import org.opensearch.client.opensearch._types.mapping.TypeMapping;
|
||||
import org.opensearch.client.opensearch._types.query_dsl.Query;
|
||||
import org.opensearch.client.opensearch._types.query_dsl.ScriptScoreQuery;
|
||||
import org.opensearch.client.opensearch.core.BulkRequest;
|
||||
import org.opensearch.client.opensearch.core.BulkResponse;
|
||||
import org.opensearch.client.opensearch.core.SearchRequest;
|
||||
import org.opensearch.client.opensearch.core.SearchResponse;
|
||||
import org.opensearch.client.opensearch.core.bulk.BulkResponseItem;
|
||||
import org.opensearch.client.transport.OpenSearchTransport;
|
||||
import org.opensearch.client.transport.endpoints.BooleanResponse;
|
||||
import org.opensearch.client.transport.httpclient5.ApacheHttpClient5TransportBuilder;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
|
||||
import javax.net.ssl.SSLContext;
|
||||
import java.io.IOException;
|
||||
import java.net.URISyntaxException;
|
||||
import java.security.KeyManagementException;
|
||||
import java.security.KeyStoreException;
|
||||
import java.security.NoSuchAlgorithmException;
|
||||
import java.util.Collection;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import static java.util.Collections.singletonList;
|
||||
import static java.util.stream.Collectors.toList;
|
||||
|
||||
/**
|
||||
* OpenSearch 向量存储
|
||||
*
|
||||
* @author songyinyin
|
||||
* @since 2024/8/10 下午8:31
|
||||
*/
|
||||
public class OpenSearchVectorStore extends DocumentStore {
|
||||
|
||||
private static final Logger log = LoggerFactory.getLogger(OpenSearchVectorStore.class);
|
||||
|
||||
private final OpenSearchClient client;
|
||||
|
||||
private final OpenSearchVectorStoreConfig config;
|
||||
|
||||
public OpenSearchVectorStore(OpenSearchVectorStoreConfig config) {
|
||||
this.config = config;
|
||||
HttpHost openSearchHost;
|
||||
try {
|
||||
openSearchHost = HttpHost.create(config.getServerUrl());
|
||||
} catch (URISyntaxException se) {
|
||||
log.error("[OpenSearch Exception]", se);
|
||||
throw new StoreException(se.getMessage());
|
||||
}
|
||||
|
||||
try {
|
||||
SSLContext sslContext = SSLContextBuilder.create().loadTrustMaterial(null, (chains, authType) -> true).build();
|
||||
TlsStrategy tlsStrategy = ClientTlsStrategyBuilder.create()
|
||||
.setSslContext(sslContext)
|
||||
.setHostnameVerifier(NoopHostnameVerifier.INSTANCE)
|
||||
.build();
|
||||
|
||||
OpenSearchTransport transport = ApacheHttpClient5TransportBuilder
|
||||
.builder(openSearchHost)
|
||||
.setMapper(new JacksonJsonpMapper())
|
||||
.setHttpClientConfigCallback(httpClientBuilder -> {
|
||||
|
||||
if (StringUtil.hasText(config.getApiKey())) {
|
||||
httpClientBuilder.setDefaultHeaders(singletonList(
|
||||
new BasicHeader("Authorization", "ApiKey " + config.getApiKey())
|
||||
));
|
||||
}
|
||||
|
||||
if (StringUtil.hasText(config.getUsername()) && StringUtil.hasText(config.getPassword())) {
|
||||
BasicCredentialsProvider credentialsProvider = new BasicCredentialsProvider();
|
||||
credentialsProvider.setCredentials(new AuthScope(openSearchHost),
|
||||
new UsernamePasswordCredentials(config.getUsername(), config.getPassword().toCharArray()));
|
||||
httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider);
|
||||
}
|
||||
|
||||
httpClientBuilder.setConnectionManager(PoolingAsyncClientConnectionManagerBuilder
|
||||
.create().setTlsStrategy(tlsStrategy).build());
|
||||
|
||||
return httpClientBuilder;
|
||||
})
|
||||
.build();
|
||||
|
||||
this.client = new OpenSearchClient(transport);
|
||||
try {
|
||||
client.ping();
|
||||
} catch (IOException e) {
|
||||
log.error("[I/O OpenSearch Exception]", e);
|
||||
throw new StoreException(e.getMessage());
|
||||
}
|
||||
} catch (NoSuchAlgorithmException | KeyManagementException | KeyStoreException e) {
|
||||
throw new StoreException("OpenSearchClient init error", e);
|
||||
}
|
||||
}
|
||||
|
||||
public OpenSearchVectorStore(OpenSearchVectorStoreConfig config, OpenSearchClient client) {
|
||||
this.config = config;
|
||||
this.client = client;
|
||||
}
|
||||
|
||||
private void createIndexIfNotExist(String indexName) {
|
||||
try {
|
||||
BooleanResponse response = client.indices().exists(c -> c.index(indexName));
|
||||
if (!response.value()) {
|
||||
log.info("[OpenSearch] Index {} not exists, creating...", indexName);
|
||||
client.indices().create(c -> c.index(indexName)
|
||||
.settings(s -> s.knn(true))
|
||||
.mappings(getDefaultMappings(this.getEmbeddingModel().dimensions())));
|
||||
}
|
||||
} catch (IOException e) {
|
||||
log.error("[I/O OpenSearch Exception]", e);
|
||||
throw new StoreException(e.getMessage());
|
||||
}
|
||||
}
|
||||
|
||||
private TypeMapping getDefaultMappings(int dimension) {
|
||||
Map<String, Property> properties = new HashMap<>(4);
|
||||
properties.put("content", Property.of(p -> p.text(TextProperty.of(t -> t))));
|
||||
properties.put("vector", Property.of(p -> p.knnVector(
|
||||
k -> k.dimension(dimension)
|
||||
)));
|
||||
return TypeMapping.of(c -> c.properties(properties));
|
||||
}
|
||||
|
||||
@Override
|
||||
public StoreResult doStore(List<Document> documents, StoreOptions options) {
|
||||
BulkRequest.Builder bulkBuilder = new BulkRequest.Builder();
|
||||
String indexName = options.getIndexNameOrDefault(config.getDefaultIndexName());
|
||||
createIndexIfNotExist(indexName);
|
||||
for (Document document : documents) {
|
||||
bulkBuilder.operations(op -> op.index(
|
||||
idx -> idx.index(indexName).id(document.getId().toString()).document(document))
|
||||
);
|
||||
}
|
||||
bulk(bulkBuilder.build());
|
||||
return StoreResult.successWithIds(documents);
|
||||
}
|
||||
|
||||
private void bulk(BulkRequest bulkRequest) {
|
||||
try {
|
||||
BulkResponse bulkResponse = client.bulk(bulkRequest);
|
||||
throwIfError(bulkResponse);
|
||||
} catch (IOException e) {
|
||||
log.error("[I/O OpenSearch Exception]", e);
|
||||
throw new StoreException(e.getMessage());
|
||||
}
|
||||
}
|
||||
|
||||
private static void throwIfError(BulkResponse bulkResponse) {
|
||||
if (bulkResponse.errors()) {
|
||||
for (BulkResponseItem item : bulkResponse.items()) {
|
||||
if (item.error() == null) {
|
||||
continue;
|
||||
}
|
||||
ErrorCause errorCause = item.error();
|
||||
throw new StoreException("type: " + errorCause.type() + "," + "reason: " + errorCause.reason());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public StoreResult doDelete(Collection<?> ids, StoreOptions options) {
|
||||
String indexName = options.getIndexNameOrDefault(config.getDefaultIndexName());
|
||||
BulkRequest.Builder bulkBuilder = new BulkRequest.Builder();
|
||||
for (Object id : ids) {
|
||||
bulkBuilder.operations(op -> op.delete(d -> d.index(indexName).id(id.toString())));
|
||||
}
|
||||
bulk(bulkBuilder.build());
|
||||
return StoreResult.success();
|
||||
}
|
||||
|
||||
@Override
|
||||
public StoreResult doUpdate(List<Document> documents, StoreOptions options) {
|
||||
BulkRequest.Builder bulkBuilder = new BulkRequest.Builder();
|
||||
String indexName = options.getIndexNameOrDefault(config.getDefaultIndexName());
|
||||
for (Document document : documents) {
|
||||
bulkBuilder.operations(op -> op.update(
|
||||
idx -> idx.index(indexName).id(document.getId().toString()).document(document))
|
||||
);
|
||||
}
|
||||
bulk(bulkBuilder.build());
|
||||
return StoreResult.successWithIds(documents);
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<Document> doSearch(SearchWrapper wrapper, StoreOptions options) {
|
||||
Double minScore = wrapper.getMinScore();
|
||||
String indexName = options.getIndexNameOrDefault(config.getDefaultIndexName());
|
||||
|
||||
// https://aws.amazon.com/cn/blogs/china/use-aws-opensearch-knn-plug-in-to-implement-vector-retrieval/
|
||||
// boost 默认是 1，小于 1 会降低相关性: https://opensearch.org/docs/latest/query-dsl/specialized/script-score/#parameters
|
||||
ScriptScoreQuery scriptScoreQuery = ScriptScoreQuery.of(q -> q.minScore(minScore == null ? 0 : minScore.floatValue())
|
||||
.query(Query.of(qu -> qu.matchAll(m -> m)))
|
||||
.script(s -> s.inline(InlineScript.of(i -> i
|
||||
.source("knn_score")
|
||||
.lang("knn")
|
||||
.params("field", JsonData.of("vector"))
|
||||
.params("query_value", JsonData.of(wrapper.getVector()))
|
||||
.params("space_type", JsonData.of("cosinesimil"))
|
||||
))));
|
||||
|
||||
try {
|
||||
SearchResponse<Document> response = client.search(
|
||||
SearchRequest.of(s -> s.index(indexName)
|
||||
.query(n -> n.scriptScore(scriptScoreQuery))
|
||||
.size(wrapper.getMaxResults())),
|
||||
Document.class
|
||||
);
|
||||
return response.hits().hits().stream()
|
||||
.filter(s -> s.source() != null)
|
||||
.map(s -> {
|
||||
Document source = s.source();
|
||||
source.setScore(s.score());
|
||||
return source;
|
||||
})
|
||||
.collect(toList());
|
||||
} catch (IOException e) {
|
||||
log.error("[I/O OpenSearch Exception]", e);
|
||||
throw new StoreException(e.getMessage());
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
/*
|
||||
* Copyright (c) 2023-2026, Easy-Agents (fuhai999@gmail.com).
|
||||
* <p>
|
||||
* 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
|
||||
* <p>
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* <p>
|
||||
* 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 com.easyagents.store.opensearch;
|
||||
|
||||
import com.easyagents.core.store.DocumentStoreConfig;
|
||||
import com.easyagents.core.util.StringUtil;
|
||||
|
||||
/**
|
||||
* 连接 open search 配置：<a href="https://opensearch.org/docs/latest/clients/java/">opensearch-java</a>
|
||||
*
|
||||
* @author songyinyin
|
||||
* @since 2024/8/10 下午8:39
|
||||
*/
|
||||
public class OpenSearchVectorStoreConfig implements DocumentStoreConfig {
|
||||
|
||||
private String serverUrl = "https://localhost:9200";
|
||||
|
||||
private String apiKey;
|
||||
|
||||
private String username;
|
||||
|
||||
private String password;
|
||||
|
||||
private String defaultIndexName = "easy-agents-default";
|
||||
|
||||
public String getServerUrl() {
|
||||
return serverUrl;
|
||||
}
|
||||
|
||||
public void setServerUrl(String serverUrl) {
|
||||
this.serverUrl = serverUrl;
|
||||
}
|
||||
|
||||
public String getApiKey() {
|
||||
return apiKey;
|
||||
}
|
||||
|
||||
public void setApiKey(String apiKey) {
|
||||
this.apiKey = apiKey;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
public String getDefaultIndexName() {
|
||||
return defaultIndexName;
|
||||
}
|
||||
|
||||
public void setDefaultIndexName(String defaultIndexName) {
|
||||
this.defaultIndexName = defaultIndexName;
|
||||
}
|
||||
|
||||
|
||||
@Override
|
||||
public boolean checkAvailable() {
|
||||
return StringUtil.hasText(this.serverUrl, this.apiKey)
|
||||
|| StringUtil.hasText(this.serverUrl, this.username, this.password);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
/*
|
||||
* Copyright (c) 2023-2026, Easy-Agents (fuhai999@gmail.com).
|
||||
* <p>
|
||||
* 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
|
||||
* <p>
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* <p>
|
||||
* 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 com.easyagents.store.opensearch;
|
||||
|
||||
import com.easyagents.core.document.Document;
|
||||
import com.easyagents.core.model.embedding.EmbeddingModel;
|
||||
import com.easyagents.core.model.embedding.EmbeddingOptions;
|
||||
import com.easyagents.core.store.SearchWrapper;
|
||||
import com.easyagents.core.store.StoreOptions;
|
||||
import com.easyagents.core.store.VectorData;
|
||||
import com.easyagents.core.store.exception.StoreException;
|
||||
import org.junit.Test;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.Arrays;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* @author songyinyin
|
||||
* @since 2024/8/11 下午3:21
|
||||
*/
|
||||
public class OpenSearchVectorStoreTest {
|
||||
|
||||
@Test(expected = StoreException.class)
|
||||
public void test01() {
|
||||
OpenSearchVectorStore store = getOpenSearchVectorStore();
|
||||
|
||||
// https://opensearch.org/docs/latest/search-plugins/vector-search/#example
|
||||
List<Document> list = new ArrayList<>();
|
||||
Document doc1 = new Document();
|
||||
doc1.setId(1);
|
||||
doc1.setContent("test1");
|
||||
doc1.setVector(new float[]{5.2f, 4.4f});
|
||||
list.add(doc1);
|
||||
Document doc2 = new Document();
|
||||
doc2.setId(2);
|
||||
doc2.setContent("test2");
|
||||
doc2.setVector(new float[]{5.2f, 3.9f});
|
||||
list.add(doc2);
|
||||
Document doc3 = new Document();
|
||||
doc3.setId(3);
|
||||
doc3.setContent("test3");
|
||||
doc3.setVector(new float[]{4.9f, 3.4f});
|
||||
list.add(doc3);
|
||||
Document doc4 = new Document();
|
||||
doc4.setId(4);
|
||||
doc4.setContent("test4");
|
||||
doc4.setVector(new float[]{4.2f, 4.6f});
|
||||
list.add(doc4);
|
||||
Document doc5 = new Document();
|
||||
doc5.setId(5);
|
||||
doc5.setContent("test5");
|
||||
doc5.setVector(new float[]{3.3f, 4.5f});
|
||||
list.add(doc5);
|
||||
store.doStore(list, StoreOptions.DEFAULT);
|
||||
|
||||
// 可能要等一会 才能查出结果
|
||||
SearchWrapper searchWrapper = new SearchWrapper();
|
||||
searchWrapper.setVector(new float[]{5, 4});
|
||||
searchWrapper.setMaxResults(3);
|
||||
List<Document> documents = store.doSearch(searchWrapper, StoreOptions.DEFAULT);
|
||||
for (Document document : documents) {
|
||||
System.out.printf("id=%s, content=%s, vector=%s, metadata=%s\n",
|
||||
document.getId(), document.getContent(), Arrays.toString(document.getVector()), document.getMetadataMap());
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
private static OpenSearchVectorStore getOpenSearchVectorStore() {
|
||||
OpenSearchVectorStoreConfig config = new OpenSearchVectorStoreConfig();
|
||||
config.setUsername("admin");
|
||||
config.setPassword("4_Pa46WQczS?");
|
||||
OpenSearchVectorStore store = new OpenSearchVectorStore(config);
|
||||
store.setEmbeddingModel(new EmbeddingModel() {
|
||||
@Override
|
||||
public VectorData embed(Document document, EmbeddingOptions options) {
|
||||
VectorData vectorData = new VectorData();
|
||||
vectorData.setVector(new float[]{0, 0});
|
||||
return vectorData;
|
||||
}
|
||||
});
|
||||
return store;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user