From 50713eac1a754e6fb1cf8af2d9d49d50486d42c1 Mon Sep 17 00:00:00 2001 From: 0007 <0007@qq.com> Date: Wed, 27 Aug 2025 19:58:19 +0800 Subject: [PATCH] Add File --- .../agentsflex/core/store/DocumentStore.java | 150 ++++++++++++++++++ 1 file changed, 150 insertions(+) create mode 100644 agents-flex-core/src/main/java/com/agentsflex/core/store/DocumentStore.java diff --git a/agents-flex-core/src/main/java/com/agentsflex/core/store/DocumentStore.java b/agents-flex-core/src/main/java/com/agentsflex/core/store/DocumentStore.java new file mode 100644 index 0000000..10b832c --- /dev/null +++ b/agents-flex-core/src/main/java/com/agentsflex/core/store/DocumentStore.java @@ -0,0 +1,150 @@ +/* + * Copyright (c) 2023-2025, Agents-Flex (fuhai999@gmail.com). + *

+ * 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 + *

+ * http://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 com.agentsflex.core.store; + +import com.agentsflex.core.llm.embedding.EmbeddingModel; +import com.agentsflex.core.document.Document; +import com.agentsflex.core.document.DocumentSplitter; +import com.agentsflex.core.document.id.DocumentIdGenerator; +import com.agentsflex.core.document.id.DocumentIdGeneratorFactory; +import com.agentsflex.core.llm.exception.LlmException; + +import java.util.Collection; +import java.util.List; + +/** + * Document Store + */ +public abstract class DocumentStore extends VectorStore { + + /** + * DocumentStore can use external embeddings models or its own embeddings + * Many vector databases come with the ability to embed themselves + */ + private EmbeddingModel embeddingModel; + + private DocumentSplitter documentSplitter; + + private DocumentIdGenerator documentIdGenerator = DocumentIdGeneratorFactory.getDocumentIdGenerator(); + + public EmbeddingModel getEmbeddingModel() { + return embeddingModel; + } + + public void setEmbeddingModel(EmbeddingModel embeddingModel) { + this.embeddingModel = embeddingModel; + } + + public DocumentSplitter getDocumentSplitter() { + return documentSplitter; + } + + public void setDocumentSplitter(DocumentSplitter documentSplitter) { + this.documentSplitter = documentSplitter; + } + + public DocumentIdGenerator getDocumentIdGenerator() { + return documentIdGenerator; + } + + public void setDocumentIdGenerator(DocumentIdGenerator documentIdGenerator) { + this.documentIdGenerator = documentIdGenerator; + } + + @Override + public StoreResult store(List documents, StoreOptions options) { + if (options == null) { + options = StoreOptions.DEFAULT; + } + + if (documentSplitter != null) { + documents = documentSplitter.splitAll(documents, documentIdGenerator); + } + // use the documentIdGenerator create unique id for document + else if (documentIdGenerator != null) { + for (Document document : documents) { + if (document.getId() == null) { + Object id = documentIdGenerator.generateId(document); + document.setId(id); + } + } + } + + embedDocumentsIfNecessary(documents, options); + + return storeInternal(documents, options); + } + + @Override + public StoreResult delete(Collection ids, StoreOptions options) { + if (options == null) { + options = StoreOptions.DEFAULT; + } + return deleteInternal(ids, options); + } + + @Override + public StoreResult update(List documents, StoreOptions options) { + if (options == null) { + options = StoreOptions.DEFAULT; + } + + embedDocumentsIfNecessary(documents, options); + return updateInternal(documents, options); + } + + + @Override + public List search(SearchWrapper wrapper, StoreOptions options) { + if (options == null) { + options = StoreOptions.DEFAULT; + } + + if (wrapper.getVector() == null && embeddingModel != null && wrapper.isWithVector()) { + VectorData vectorData = embeddingModel.embed(Document.of(wrapper.getText()), options.getEmbeddingOptions()); + if (vectorData == null) { + throw new LlmException("Embedding model does not contain vector data"); + } + wrapper.setVector(vectorData.getVector()); + } + + return searchInternal(wrapper, options); + } + + + protected void embedDocumentsIfNecessary(List documents, StoreOptions options) { + if (embeddingModel == null) { + return; + } + for (Document document : documents) { + if (document.getVector() == null) { + VectorData vectorData = embeddingModel.embed(document, options.getEmbeddingOptions()); + if (vectorData != null) { + document.setVector(vectorData.getVector()); + } + } + } + } + + + public abstract StoreResult storeInternal(List documents, StoreOptions options); + + public abstract StoreResult deleteInternal(Collection ids, StoreOptions options); + + public abstract StoreResult updateInternal(List documents, StoreOptions options); + + public abstract List searchInternal(SearchWrapper wrapper, StoreOptions options); +}