001/*
002 *  Copyright (c) 2023-2026, Agents-Flex (fuhai999@gmail.com).
003 *  <p>
004 *  Licensed under the Apache License, Version 2.0 (the "License");
005 *  you may not use this file except in compliance with the License.
006 *  You may obtain a copy of the License at
007 *  <p>
008 *  http://www.apache.org/licenses/LICENSE-2.0
009 *  <p>
010 *  Unless required by applicable law or agreed to in writing, software
011 *  distributed under the License is distributed on an "AS IS" BASIS,
012 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 *  See the License for the specific language governing permissions and
014 *  limitations under the License.
015 */
016package com.agentsflex.core.store;
017
018import com.agentsflex.core.model.embedding.EmbeddingModel;
019import com.agentsflex.core.document.Document;
020import com.agentsflex.core.document.DocumentSplitter;
021import com.agentsflex.core.document.id.DocumentIdGenerator;
022import com.agentsflex.core.document.id.DocumentIdGeneratorFactory;
023import com.agentsflex.core.model.exception.ModelException;
024
025import java.util.Collection;
026import java.util.List;
027
028/**
029 * Document Store
030 */
031public abstract class DocumentStore extends VectorStore<Document> {
032
033    /**
034     * DocumentStore can use external embeddings models or its own embeddings
035     * Many vector databases come with the ability to embed themselves
036     */
037    private EmbeddingModel embeddingModel;
038
039    private DocumentSplitter documentSplitter;
040
041    private DocumentIdGenerator documentIdGenerator = DocumentIdGeneratorFactory.getDocumentIdGenerator();
042
043    public EmbeddingModel getEmbeddingModel() {
044        return embeddingModel;
045    }
046
047    public void setEmbeddingModel(EmbeddingModel embeddingModel) {
048        this.embeddingModel = embeddingModel;
049    }
050
051    public DocumentSplitter getDocumentSplitter() {
052        return documentSplitter;
053    }
054
055    public void setDocumentSplitter(DocumentSplitter documentSplitter) {
056        this.documentSplitter = documentSplitter;
057    }
058
059    public DocumentIdGenerator getDocumentIdGenerator() {
060        return documentIdGenerator;
061    }
062
063    public void setDocumentIdGenerator(DocumentIdGenerator documentIdGenerator) {
064        this.documentIdGenerator = documentIdGenerator;
065    }
066
067    @Override
068    public StoreResult store(List<Document> documents, StoreOptions options) {
069        if (options == null) {
070            options = StoreOptions.DEFAULT;
071        }
072
073        if (documentSplitter != null) {
074            documents = documentSplitter.splitAll(documents, documentIdGenerator);
075        }
076        // use the documentIdGenerator create unique id for document
077        else if (documentIdGenerator != null) {
078            for (Document document : documents) {
079                if (document.getId() == null) {
080                    Object id = documentIdGenerator.generateId(document);
081                    document.setId(id);
082                }
083            }
084        }
085
086        embedDocumentsIfNecessary(documents, options);
087
088        return doStore(documents, options);
089    }
090
091    @Override
092    public StoreResult delete(Collection<?> ids, StoreOptions options) {
093        if (options == null) {
094            options = StoreOptions.DEFAULT;
095        }
096        return doDelete(ids, options);
097    }
098
099    @Override
100    public StoreResult update(List<Document> documents, StoreOptions options) {
101        if (options == null) {
102            options = StoreOptions.DEFAULT;
103        }
104
105        embedDocumentsIfNecessary(documents, options);
106        return doUpdate(documents, options);
107    }
108
109
110    @Override
111    public List<Document> search(SearchWrapper wrapper, StoreOptions options) {
112        if (options == null) {
113            options = StoreOptions.DEFAULT;
114        }
115
116        if (wrapper.getVector() == null && embeddingModel != null && wrapper.isWithVector()) {
117            VectorData vectorData = embeddingModel.embed(Document.of(wrapper.getText()), options.getEmbeddingOptions());
118            if (vectorData == null) {
119                throw new ModelException("Embedding model does not contain vector data");
120            }
121            wrapper.setVector(vectorData.getVector());
122        }
123
124        return doSearch(wrapper, options);
125    }
126
127
128    protected void embedDocumentsIfNecessary(List<Document> documents, StoreOptions options) {
129        if (embeddingModel == null) {
130            return;
131        }
132        for (Document document : documents) {
133            if (document.getVector() == null) {
134                VectorData vectorData = embeddingModel.embed(document, options.getEmbeddingOptions());
135                if (vectorData != null) {
136                    document.setVector(vectorData.getVector());
137                }
138            }
139        }
140    }
141
142
143    public abstract StoreResult doStore(List<Document> documents, StoreOptions options);
144
145    public abstract StoreResult doDelete(Collection<?> ids, StoreOptions options);
146
147    public abstract StoreResult doUpdate(List<Document> documents, StoreOptions options);
148
149    public abstract List<Document> doSearch(SearchWrapper wrapper, StoreOptions options);
150}