content: /*
   *  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.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<Document> {
  
      /**
       * 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<Document> 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<Document> documents, StoreOptions options) {
          if (options == null) {
              options = StoreOptions.DEFAULT;
          }
  
          embedDocumentsIfNecessary(documents, options);
          return updateInternal(documents, options);
      }
  
  
      @Override
      public List<Document> 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<Document> 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<Document> documents, StoreOptions options);
  
      public abstract StoreResult deleteInternal(Collection<?> ids, StoreOptions options);
  
      public abstract StoreResult updateInternal(List<Document> documents, StoreOptions options);
  
      public abstract List<Document> searchInternal(SearchWrapper wrapper, StoreOptions options);
  }
  