mirror of
https://gitee.com/agents-flex/agents-flex.git
synced 2024-12-03 12:27:53 +08:00
!11 add elasticsearch vector store and springboot auto config
Merge pull request !11 from 读钓/main
This commit is contained in:
commit
24e8091e2d
@ -103,6 +103,18 @@
|
|||||||
<version>${agents-flex.version}</version>
|
<version>${agents-flex.version}</version>
|
||||||
<optional>true</optional>
|
<optional>true</optional>
|
||||||
</dependency>
|
</dependency>
|
||||||
|
<dependency>
|
||||||
|
<groupId>com.agentsflex</groupId>
|
||||||
|
<artifactId>agents-flex-store-elasticsearch</artifactId>
|
||||||
|
<version>${agents-flex.version}</version>
|
||||||
|
<optional>true</optional>
|
||||||
|
</dependency>
|
||||||
|
<dependency>
|
||||||
|
<groupId>com.agentsflex</groupId>
|
||||||
|
<artifactId>agents-flex-store-opensearch</artifactId>
|
||||||
|
<version>${agents-flex.version}</version>
|
||||||
|
<optional>true</optional>
|
||||||
|
</dependency>
|
||||||
<!--store end-->
|
<!--store end-->
|
||||||
</dependencies>
|
</dependencies>
|
||||||
|
|
||||||
|
@ -0,0 +1,52 @@
|
|||||||
|
/*
|
||||||
|
* Copyright (c) 2023-2025, Agents-Flex (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.agentsflex.spring.boot.store.elasticsearch;
|
||||||
|
|
||||||
|
import co.elastic.clients.elasticsearch.ElasticsearchClient;
|
||||||
|
import com.agentsflex.store.elasticsearch.ElasticSearchVectorStore;
|
||||||
|
import com.agentsflex.store.elasticsearch.ElasticSearchVectorStoreConfig;
|
||||||
|
import org.springframework.beans.factory.annotation.Autowired;
|
||||||
|
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.context.annotation.Configuration;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* @author songyinyin
|
||||||
|
* @since 2024/8/13 上午11:26
|
||||||
|
*/
|
||||||
|
@Configuration(proxyBeanMethods = false)
|
||||||
|
@ConditionalOnClass(ElasticSearchVectorStore.class)
|
||||||
|
@EnableConfigurationProperties(ElasticSearchProperties.class)
|
||||||
|
public class ElasticSearchAutoConfiguration {
|
||||||
|
|
||||||
|
@Bean
|
||||||
|
@ConditionalOnMissingBean
|
||||||
|
public ElasticSearchVectorStore elasticSearchVectorStore(ElasticSearchProperties properties,
|
||||||
|
@Autowired(required = false) ElasticsearchClient client) {
|
||||||
|
ElasticSearchVectorStoreConfig config = new ElasticSearchVectorStoreConfig();
|
||||||
|
config.setServerUrl(properties.getServerUrl());
|
||||||
|
config.setApiKey(properties.getApiKey());
|
||||||
|
config.setUsername(properties.getUsername());
|
||||||
|
config.setPassword(properties.getPassword());
|
||||||
|
config.setDefaultIndexName(properties.getDefaultIndexName());
|
||||||
|
if (client != null) {
|
||||||
|
return new ElasticSearchVectorStore(config, client);
|
||||||
|
}
|
||||||
|
return new ElasticSearchVectorStore(config);
|
||||||
|
}
|
||||||
|
}
|
@ -0,0 +1,76 @@
|
|||||||
|
/*
|
||||||
|
* Copyright (c) 2023-2025, Agents-Flex (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.agentsflex.spring.boot.store.elasticsearch;
|
||||||
|
|
||||||
|
import org.springframework.boot.context.properties.ConfigurationProperties;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* @author songyinyin
|
||||||
|
* @since 2024/8/13 上午11:25
|
||||||
|
*/
|
||||||
|
@ConfigurationProperties(prefix = "agents-flex.store.elasticsearch")
|
||||||
|
public class ElasticSearchProperties {
|
||||||
|
|
||||||
|
private String serverUrl = "https://localhost:9200";
|
||||||
|
|
||||||
|
private String apiKey;
|
||||||
|
|
||||||
|
private String username;
|
||||||
|
|
||||||
|
private String password;
|
||||||
|
|
||||||
|
private String defaultIndexName = "agents-flex-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;
|
||||||
|
}
|
||||||
|
}
|
@ -0,0 +1,52 @@
|
|||||||
|
/*
|
||||||
|
* Copyright (c) 2023-2025, Agents-Flex (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.agentsflex.spring.boot.store.opensearch;
|
||||||
|
|
||||||
|
import com.agentsflex.store.opensearch.OpenSearchVectorStore;
|
||||||
|
import com.agentsflex.store.opensearch.OpenSearchVectorStoreConfig;
|
||||||
|
import org.opensearch.client.opensearch.OpenSearchClient;
|
||||||
|
import org.springframework.beans.factory.annotation.Autowired;
|
||||||
|
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.context.annotation.Configuration;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* @author songyinyin
|
||||||
|
* @since 2024/8/13 上午11:26
|
||||||
|
*/
|
||||||
|
@Configuration(proxyBeanMethods = false)
|
||||||
|
@ConditionalOnClass(OpenSearchVectorStore.class)
|
||||||
|
@EnableConfigurationProperties(OpenSearchProperties.class)
|
||||||
|
public class OpenSearchAutoConfiguration {
|
||||||
|
|
||||||
|
@Bean
|
||||||
|
@ConditionalOnMissingBean
|
||||||
|
public OpenSearchVectorStore openSearchVectorStore(OpenSearchProperties properties,
|
||||||
|
@Autowired(required = false) OpenSearchClient client) {
|
||||||
|
OpenSearchVectorStoreConfig config = new OpenSearchVectorStoreConfig();
|
||||||
|
config.setServerUrl(properties.getServerUrl());
|
||||||
|
config.setApiKey(properties.getApiKey());
|
||||||
|
config.setUsername(properties.getUsername());
|
||||||
|
config.setPassword(properties.getPassword());
|
||||||
|
config.setDefaultIndexName(properties.getDefaultIndexName());
|
||||||
|
if (client != null) {
|
||||||
|
return new OpenSearchVectorStore(config, client);
|
||||||
|
}
|
||||||
|
return new OpenSearchVectorStore(config);
|
||||||
|
}
|
||||||
|
}
|
@ -0,0 +1,76 @@
|
|||||||
|
/*
|
||||||
|
* Copyright (c) 2023-2025, Agents-Flex (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.agentsflex.spring.boot.store.opensearch;
|
||||||
|
|
||||||
|
import org.springframework.boot.context.properties.ConfigurationProperties;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* @author songyinyin
|
||||||
|
* @since 2024/8/13 上午11:25
|
||||||
|
*/
|
||||||
|
@ConfigurationProperties(prefix = "agents-flex.store.opensearch")
|
||||||
|
public class OpenSearchProperties {
|
||||||
|
|
||||||
|
private String serverUrl = "https://localhost:9200";
|
||||||
|
|
||||||
|
private String apiKey;
|
||||||
|
|
||||||
|
private String username;
|
||||||
|
|
||||||
|
private String password;
|
||||||
|
|
||||||
|
private String defaultIndexName = "agents-flex-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;
|
||||||
|
}
|
||||||
|
}
|
@ -12,9 +12,38 @@
|
|||||||
<artifactId>agents-flex-store-elasticsearch</artifactId>
|
<artifactId>agents-flex-store-elasticsearch</artifactId>
|
||||||
|
|
||||||
<properties>
|
<properties>
|
||||||
|
<elasticsearch.version>8.15.0</elasticsearch.version>
|
||||||
|
<jackson.version>2.17.0</jackson.version>
|
||||||
|
|
||||||
<maven.compiler.source>8</maven.compiler.source>
|
<maven.compiler.source>8</maven.compiler.source>
|
||||||
<maven.compiler.target>8</maven.compiler.target>
|
<maven.compiler.target>8</maven.compiler.target>
|
||||||
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
||||||
</properties>
|
</properties>
|
||||||
|
|
||||||
|
<dependencies>
|
||||||
|
<dependency>
|
||||||
|
<groupId>com.agentsflex</groupId>
|
||||||
|
<artifactId>agents-flex-core</artifactId>
|
||||||
|
<version>1.0.0-beta.8</version>
|
||||||
|
</dependency>
|
||||||
|
|
||||||
|
<dependency>
|
||||||
|
<groupId>co.elastic.clients</groupId>
|
||||||
|
<artifactId>elasticsearch-java</artifactId>
|
||||||
|
<version>${elasticsearch.version}</version>
|
||||||
|
</dependency>
|
||||||
|
|
||||||
|
<dependency>
|
||||||
|
<groupId>com.fasterxml.jackson.core</groupId>
|
||||||
|
<artifactId>jackson-databind</artifactId>
|
||||||
|
<version>${jackson.version}</version>
|
||||||
|
</dependency>
|
||||||
|
|
||||||
|
<dependency>
|
||||||
|
<groupId>junit</groupId>
|
||||||
|
<artifactId>junit</artifactId>
|
||||||
|
<scope>test</scope>
|
||||||
|
</dependency>
|
||||||
|
</dependencies>
|
||||||
|
|
||||||
</project>
|
</project>
|
||||||
|
@ -1,7 +0,0 @@
|
|||||||
package com.agentsflex;
|
|
||||||
|
|
||||||
public class Main {
|
|
||||||
public static void main(String[] args) {
|
|
||||||
System.out.println("Hello world!");
|
|
||||||
}
|
|
||||||
}
|
|
@ -0,0 +1,237 @@
|
|||||||
|
/*
|
||||||
|
* Copyright (c) 2023-2025, Agents-Flex (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.agentsflex.store.elasticsearch;
|
||||||
|
|
||||||
|
import co.elastic.clients.elasticsearch.ElasticsearchClient;
|
||||||
|
import co.elastic.clients.elasticsearch._types.ErrorCause;
|
||||||
|
import co.elastic.clients.elasticsearch._types.mapping.DenseVectorProperty;
|
||||||
|
import co.elastic.clients.elasticsearch._types.mapping.Property;
|
||||||
|
import co.elastic.clients.elasticsearch._types.mapping.TextProperty;
|
||||||
|
import co.elastic.clients.elasticsearch._types.mapping.TypeMapping;
|
||||||
|
import co.elastic.clients.elasticsearch._types.query_dsl.Query;
|
||||||
|
import co.elastic.clients.elasticsearch._types.query_dsl.ScriptScoreQuery;
|
||||||
|
import co.elastic.clients.elasticsearch.core.BulkRequest;
|
||||||
|
import co.elastic.clients.elasticsearch.core.BulkResponse;
|
||||||
|
import co.elastic.clients.elasticsearch.core.SearchRequest;
|
||||||
|
import co.elastic.clients.elasticsearch.core.SearchResponse;
|
||||||
|
import co.elastic.clients.elasticsearch.core.bulk.BulkResponseItem;
|
||||||
|
import co.elastic.clients.json.JsonData;
|
||||||
|
import co.elastic.clients.json.jackson.JacksonJsonpMapper;
|
||||||
|
import co.elastic.clients.transport.ElasticsearchTransport;
|
||||||
|
import co.elastic.clients.transport.endpoints.BooleanResponse;
|
||||||
|
import co.elastic.clients.transport.rest_client.RestClientTransport;
|
||||||
|
import com.agentsflex.core.document.Document;
|
||||||
|
import com.agentsflex.core.store.DocumentStore;
|
||||||
|
import com.agentsflex.core.store.SearchWrapper;
|
||||||
|
import com.agentsflex.core.store.StoreOptions;
|
||||||
|
import com.agentsflex.core.store.StoreResult;
|
||||||
|
import com.agentsflex.core.store.exception.StoreException;
|
||||||
|
import com.agentsflex.core.util.StringUtil;
|
||||||
|
import org.apache.http.Header;
|
||||||
|
import org.apache.http.HttpHost;
|
||||||
|
import org.apache.http.auth.AuthScope;
|
||||||
|
import org.apache.http.auth.UsernamePasswordCredentials;
|
||||||
|
import org.apache.http.client.CredentialsProvider;
|
||||||
|
import org.apache.http.impl.client.BasicCredentialsProvider;
|
||||||
|
import org.apache.http.message.BasicHeader;
|
||||||
|
import org.apache.http.ssl.SSLContextBuilder;
|
||||||
|
import org.elasticsearch.client.RestClient;
|
||||||
|
import org.elasticsearch.client.RestClientBuilder;
|
||||||
|
import org.slf4j.Logger;
|
||||||
|
import org.slf4j.LoggerFactory;
|
||||||
|
|
||||||
|
import javax.net.ssl.SSLContext;
|
||||||
|
import java.io.IOException;
|
||||||
|
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.stream.Collectors.toList;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* es 向量存储:<a href="https://www.elastic.co/guide/en/elasticsearch/client/java-api-client/current/introduction.html">elasticsearch-java</a>
|
||||||
|
*
|
||||||
|
* @author songyinyin
|
||||||
|
* @since 2024/8/12 下午4:17
|
||||||
|
*/
|
||||||
|
public class ElasticSearchVectorStore extends DocumentStore {
|
||||||
|
|
||||||
|
private static final Logger log = LoggerFactory.getLogger(ElasticSearchVectorStore.class);
|
||||||
|
|
||||||
|
private final ElasticsearchClient client;
|
||||||
|
|
||||||
|
private final ElasticSearchVectorStoreConfig config;
|
||||||
|
|
||||||
|
public ElasticSearchVectorStore(ElasticSearchVectorStoreConfig config) {
|
||||||
|
this.config = config;
|
||||||
|
RestClientBuilder restClientBuilder = RestClient.builder(HttpHost.create(config.getServerUrl()));
|
||||||
|
|
||||||
|
try {
|
||||||
|
SSLContext sslContext = SSLContextBuilder.create().loadTrustMaterial(null, (chains, authType) -> true).build();
|
||||||
|
|
||||||
|
if (StringUtil.hasText(config.getUsername())) {
|
||||||
|
CredentialsProvider provider = new BasicCredentialsProvider();
|
||||||
|
provider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(config.getUsername(), config.getPassword()));
|
||||||
|
restClientBuilder.setHttpClientConfigCallback(httpClientBuilder -> {
|
||||||
|
httpClientBuilder.setSSLContext(sslContext);
|
||||||
|
httpClientBuilder.setDefaultCredentialsProvider(provider);
|
||||||
|
return httpClientBuilder;
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
if (StringUtil.hasText(config.getApiKey())) {
|
||||||
|
restClientBuilder.setDefaultHeaders(new Header[]{
|
||||||
|
new BasicHeader("Authorization", "Apikey " + config.getApiKey())
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
ElasticsearchTransport transport = new RestClientTransport(restClientBuilder.build(), new JacksonJsonpMapper());
|
||||||
|
|
||||||
|
this.client = new ElasticsearchClient(transport);
|
||||||
|
} catch (NoSuchAlgorithmException | KeyManagementException | KeyStoreException e) {
|
||||||
|
throw new StoreException("Elasticsearch init error", e);
|
||||||
|
}
|
||||||
|
try {
|
||||||
|
client.ping();
|
||||||
|
} catch (IOException e) {
|
||||||
|
log.error("[I/O Elasticsearch Exception]", e);
|
||||||
|
throw new StoreException(e.getMessage());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
public ElasticSearchVectorStore(ElasticSearchVectorStoreConfig config, ElasticsearchClient client) {
|
||||||
|
this.config = config;
|
||||||
|
this.client = client;
|
||||||
|
}
|
||||||
|
|
||||||
|
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 storeInternal(List<Document> documents, StoreOptions options) {
|
||||||
|
String indexName = options.getIndexNameOrDefault(config.getDefaultIndexName());
|
||||||
|
createIndexIfNotExist(indexName);
|
||||||
|
return saveOrUpdate(documents, indexName);
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public StoreResult deleteInternal(Collection<Object> 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 updateInternal(List<Document> documents, StoreOptions options) {
|
||||||
|
String indexName = options.getIndexNameOrDefault(config.getDefaultIndexName());
|
||||||
|
return saveOrUpdate(documents, indexName);
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public List<Document> searchInternal(SearchWrapper wrapper, StoreOptions options) {
|
||||||
|
Double minScore = wrapper.getMinScore();
|
||||||
|
String indexName = options.getIndexNameOrDefault(config.getDefaultIndexName());
|
||||||
|
|
||||||
|
// https://www.elastic.co/guide/en/elasticsearch/reference/current/knn-search.html
|
||||||
|
ScriptScoreQuery scriptScoreQuery = ScriptScoreQuery.of(fn -> fn
|
||||||
|
.minScore(minScore == null ? 0 : minScore.floatValue())
|
||||||
|
.query(Query.of(q -> q.matchAll(m -> m)))
|
||||||
|
.script(s -> s
|
||||||
|
.source("(cosineSimilarity(params.query_vector, 'vector') + 1.0) / 2")
|
||||||
|
.params("query_vector", JsonData.of(wrapper.getVector()))
|
||||||
|
)
|
||||||
|
);
|
||||||
|
|
||||||
|
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.addMetadata("_score", s.score());
|
||||||
|
return source;
|
||||||
|
})
|
||||||
|
.collect(toList());
|
||||||
|
} catch (IOException e) {
|
||||||
|
log.error("[I/O Elasticsearch Exception]", e);
|
||||||
|
throw new StoreException(e.getMessage());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private StoreResult saveOrUpdate(List<Document> documents, String indexName) {
|
||||||
|
BulkRequest.Builder bulkBuilder = new BulkRequest.Builder();
|
||||||
|
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 Elasticsearch Exception]", e);
|
||||||
|
throw new StoreException(e.getMessage());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private void createIndexIfNotExist(String indexName) {
|
||||||
|
try {
|
||||||
|
BooleanResponse response = client.indices().exists(c -> c.index(indexName));
|
||||||
|
if (!response.value()) {
|
||||||
|
log.info("[ElasticSearch] Index {} not exists, creating...", indexName);
|
||||||
|
client.indices().create(c -> c.index(indexName)
|
||||||
|
.mappings(getDefaultMappings(this.getEmbeddingModel().dimensions())));
|
||||||
|
}
|
||||||
|
} catch (IOException e) {
|
||||||
|
log.error("[I/O ElasticSearch 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.denseVector(DenseVectorProperty.of(d -> d.dims(dimension)))));
|
||||||
|
return TypeMapping.of(c -> c.properties(properties));
|
||||||
|
}
|
||||||
|
}
|
@ -0,0 +1,76 @@
|
|||||||
|
/*
|
||||||
|
* Copyright (c) 2023-2025, Agents-Flex (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.agentsflex.store.elasticsearch;
|
||||||
|
|
||||||
|
import java.io.Serializable;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 连接 elasticsearch 配置:<a href="https://www.elastic.co/guide/en/elasticsearch/client/java-api-client/current/getting-started-java.html">elasticsearch-java</a>
|
||||||
|
*
|
||||||
|
* @author songyinyin
|
||||||
|
*/
|
||||||
|
public class ElasticSearchVectorStoreConfig implements Serializable {
|
||||||
|
|
||||||
|
private String serverUrl = "https://localhost:9200";
|
||||||
|
|
||||||
|
private String apiKey;
|
||||||
|
|
||||||
|
private String username;
|
||||||
|
|
||||||
|
private String password;
|
||||||
|
|
||||||
|
private String defaultIndexName = "agents-flex-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;
|
||||||
|
}
|
||||||
|
}
|
@ -0,0 +1,98 @@
|
|||||||
|
/*
|
||||||
|
* Copyright (c) 2023-2025, Agents-Flex (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.agentsflex.store.opensearch;
|
||||||
|
|
||||||
|
import com.agentsflex.core.document.Document;
|
||||||
|
import com.agentsflex.core.llm.embedding.EmbeddingModel;
|
||||||
|
import com.agentsflex.core.llm.embedding.EmbeddingOptions;
|
||||||
|
import com.agentsflex.core.store.SearchWrapper;
|
||||||
|
import com.agentsflex.core.store.StoreOptions;
|
||||||
|
import com.agentsflex.core.store.VectorData;
|
||||||
|
import com.agentsflex.store.elasticsearch.ElasticSearchVectorStore;
|
||||||
|
import com.agentsflex.store.elasticsearch.ElasticSearchVectorStoreConfig;
|
||||||
|
import org.junit.Test;
|
||||||
|
|
||||||
|
import java.util.ArrayList;
|
||||||
|
import java.util.Arrays;
|
||||||
|
import java.util.List;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* @author songyinyin
|
||||||
|
*/
|
||||||
|
public class ElasticSearchVectorStoreTest {
|
||||||
|
|
||||||
|
private static ElasticSearchVectorStore getVectorStore() {
|
||||||
|
ElasticSearchVectorStoreConfig config = new ElasticSearchVectorStoreConfig();
|
||||||
|
// config.setApiKey("bmtXRVNaRUJNMEZXZzMzcnNvSXk6MlNMVmFnT0hRVVNUSmN3UXpoNWp4Zw==");
|
||||||
|
config.setUsername("elastic");
|
||||||
|
config.setPassword("Dd2024a10");
|
||||||
|
ElasticSearchVectorStore store = new ElasticSearchVectorStore(config);
|
||||||
|
store.setEmbeddingModel(new EmbeddingModel() {
|
||||||
|
@Override
|
||||||
|
public VectorData embed(Document document, EmbeddingOptions options) {
|
||||||
|
VectorData vectorData = new VectorData();
|
||||||
|
vectorData.setVector(new double[]{0, 0});
|
||||||
|
return vectorData;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
return store;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
public void test01() {
|
||||||
|
ElasticSearchVectorStore store = getVectorStore();
|
||||||
|
|
||||||
|
// 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 double[]{5.2, 4.4});
|
||||||
|
list.add(doc1);
|
||||||
|
Document doc2 = new Document();
|
||||||
|
doc2.setId(2);
|
||||||
|
doc2.setContent("test2");
|
||||||
|
doc2.setVector(new double[]{5.2, 3.9});
|
||||||
|
list.add(doc2);
|
||||||
|
Document doc3 = new Document();
|
||||||
|
doc3.setId(3);
|
||||||
|
doc3.setContent("test3");
|
||||||
|
doc3.setVector(new double[]{4.9, 3.4});
|
||||||
|
list.add(doc3);
|
||||||
|
Document doc4 = new Document();
|
||||||
|
doc4.setId(4);
|
||||||
|
doc4.setContent("test4");
|
||||||
|
doc4.setVector(new double[]{4.2, 4.6});
|
||||||
|
list.add(doc4);
|
||||||
|
Document doc5 = new Document();
|
||||||
|
doc5.setId(5);
|
||||||
|
doc5.setContent("test5");
|
||||||
|
doc5.setVector(new double[]{3.3, 4.5});
|
||||||
|
list.add(doc5);
|
||||||
|
store.storeInternal(list, StoreOptions.DEFAULT);
|
||||||
|
|
||||||
|
// 可能要等一会 才能查出结果
|
||||||
|
SearchWrapper searchWrapper = new SearchWrapper();
|
||||||
|
searchWrapper.setVector(new double[]{5, 4});
|
||||||
|
searchWrapper.setMaxResults(3);
|
||||||
|
List<Document> documents = store.searchInternal(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());
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
}
|
@ -135,6 +135,11 @@ public class OpenSearchVectorStore extends DocumentStore {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
public OpenSearchVectorStore(OpenSearchVectorStoreConfig config, OpenSearchClient client) {
|
||||||
|
this.config = config;
|
||||||
|
this.client = client;
|
||||||
|
}
|
||||||
|
|
||||||
private void createIndexIfNotExist(String indexName) {
|
private void createIndexIfNotExist(String indexName) {
|
||||||
try {
|
try {
|
||||||
BooleanResponse response = client.indices().exists(c -> c.index(indexName));
|
BooleanResponse response = client.indices().exists(c -> c.index(indexName));
|
||||||
|
Loading…
Reference in New Issue
Block a user