Source code

001/*
002 *  Copyright 2026 Anyware Services
003 *
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 *
008 *      http://www.apache.org/licenses/LICENSE-2.0
009 *
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 org.ametys.plugins.ai.search;
017
018import java.io.IOException;
019import java.io.InputStreamReader;
020import java.io.Reader;
021import java.time.ZonedDateTime;
022import java.util.ArrayList;
023import java.util.Arrays;
024import java.util.Collection;
025import java.util.HashMap;
026import java.util.List;
027import java.util.Map;
028import java.util.Set;
029
030import javax.jcr.LoginException;
031import javax.jcr.RepositoryException;
032
033import org.apache.avalon.framework.context.Context;
034import org.apache.avalon.framework.context.ContextException;
035import org.apache.avalon.framework.context.Contextualizable;
036import org.apache.avalon.framework.service.ServiceException;
037import org.apache.avalon.framework.service.ServiceManager;
038import org.apache.cocoon.components.ContextHelper;
039import org.apache.cocoon.components.source.impl.SitemapSource;
040import org.apache.cocoon.environment.Request;
041import org.apache.commons.io.IOUtils;
042import org.apache.commons.lang3.StringUtils;
043import org.apache.commons.lang3.Strings;
044import org.apache.excalibur.source.SourceResolver;
045import org.apache.solr.common.SolrInputDocument;
046
047import org.ametys.cms.content.indexing.solr.SolrFieldNames;
048import org.ametys.cms.content.indexing.solr.SolrIndexer;
049import org.ametys.cms.repository.Content;
050import org.ametys.cms.search.query.AndQuery;
051import org.ametys.cms.search.query.DocumentTypeQuery;
052import org.ametys.cms.search.query.Query;
053import org.ametys.cms.search.solr.schema.FieldDefinition;
054import org.ametys.cms.search.solr.schema.FieldTypeDefinition;
055import org.ametys.cms.search.solr.schema.SchemaDefinition;
056import org.ametys.core.authentication.AuthenticateAction;
057import org.ametys.core.user.UserIdentity;
058import org.ametys.core.user.population.UserPopulationDAO;
059import org.ametys.core.util.URIUtils;
060import org.ametys.plugins.ai.AIHelper;
061import org.ametys.plugins.ai.provider.impl.DefaultTokenCountEstimator;
062import org.ametys.plugins.repository.AmetysObject;
063import org.ametys.plugins.repository.AmetysRepositoryException;
064import org.ametys.plugins.repository.jcr.DefaultTraversableAmetysObject;
065import org.ametys.plugins.repository.provider.RequestAttributeWorkspaceSelector;
066import org.ametys.plugins.repository.version.VersionableAmetysObject;
067import org.ametys.runtime.plugin.component.AbstractLogEnabled;
068import org.ametys.runtime.plugin.component.DeferredServiceable;
069import org.ametys.web.WebConstants;
070import org.ametys.web.indexing.solr.AdditionalWebDataIndexer;
071import org.ametys.web.indexing.solr.SolrWebFieldNames;
072import org.ametys.web.repository.content.WebContent;
073import org.ametys.web.search.query.SiteQuery;
074import org.ametys.web.search.query.SitemapQuery;
075
076import dev.langchain4j.data.document.Document;
077import dev.langchain4j.data.document.DocumentSplitter;
078import dev.langchain4j.data.document.splitter.DocumentSplitters;
079import dev.langchain4j.data.segment.TextSegment;
080import dev.langchain4j.model.TokenCountEstimator;
081
082/**
083 * Indexer for AI additional data.
084 */
085public class AIAdditionalDataIndexer extends AbstractLogEnabled implements AdditionalWebDataIndexer, DeferredServiceable, Contextualizable
086{
087    /** The solr type for chunk embedding */
088    public static final String SOLR_TYPE = "chunk_embedding";
089    /** The solr field name for embedding vector */
090    public static final String SOLR_EMBEDDING = "chunk_embedding";
091    /** The solr field name with the content id */
092    public static final String SOLR_CONTENT_ID = "chunk_object_s_dv"; // s_dv required for Ametys join queries
093    
094    /** The content attribute name for a chunk */
095    public static final String DATA_CHUNK_CONTENT = "content";
096    /** The embedding attribute name for a chunk */
097    public static final String DATA_CHUNK_EMBDEDING = "embedding";
098    /** The creation datetime attribute name for a chunk */
099    public static final String DATA_CHUNK_CREATION_DATETIME = "creation";
100    /** The identifier of the embedding used */
101    public static final String DATA_CHUNK_CREATOR = "creator";
102
103
104    private static final int __DOCUMENT_CHUNK_MAX_SIZE = 800;
105    private static final int __DOCUMENT_CHUNK_MAX_OVERLAP = 400;
106
107    private AIHelper _aiHelper;
108
109    private SourceResolver _sourceResolver;
110    private Context _context;
111    private CacheDAO _cacheDAO;
112    private SolrIndexer _solrIndexer;
113
114    public void deferredService(ServiceManager manager) throws ServiceException
115    {
116        _aiHelper = (AIHelper) manager.lookup(AIHelper.ROLE);
117        _sourceResolver = (SourceResolver) manager.lookup(SourceResolver.ROLE);
118        _cacheDAO = (CacheDAO) manager.lookup(CacheDAO.ROLE);
119        _solrIndexer = (SolrIndexer) manager.lookup(SolrIndexer.ROLE);
120    }
121    
122    public void contextualize(Context context) throws ContextException
123    {
124        _context = context;
125    }
126    
127    public Collection<String> getSupportedTypes()
128    {
129        return Set.of("content");
130    }
131    
132    public List<SolrInputDocument> indexAdditionalDocuments(AmetysObject object, SolrInputDocument document)
133    {
134        if (!(object instanceof Content content))
135        {
136            throw new IllegalArgumentException("Object should be a Content and not a " + object.getClass().getName());
137        }
138        
139        if (!_aiHelper.isEmbeddingSupported()
140                || content instanceof WebContent webContent && !_aiHelper.isSemanticIndexationActivated(webContent.getSiteName()))
141        {
142            return List.of();
143        }
144        
145        try
146        {
147            DefaultTraversableAmetysObject cacheContentVersionNode = _updateAndGetCache(content);
148            
149            return _getDocumentsFromCache(content, cacheContentVersionNode);
150        }
151        catch (IOException | AmetysRepositoryException | RepositoryException e)
152        {
153            throw new RuntimeException("An error occurred while preparing content " + content.getId(), e);
154        }
155    }
156    
157    private DefaultTraversableAmetysObject _updateAndGetCache(Content content) throws IOException, AmetysRepositoryException, LoginException, RepositoryException
158    {
159        String embeddingGeneratorKey = _aiHelper.getEmbeddingGeneratorKey();
160        
161        // Convert to md and clean anchors
162        List<TextSegment> segments = _contentToSegments(content);
163        String revision = _getContentRevision(content);
164        
165        DefaultTraversableAmetysObject nearestCacheContentVersionNode = _cacheDAO.getNearestCacheContentVersionNode(content, revision);
166        
167        DefaultTraversableAmetysObject cacheContentVersionNode = _checkCacheValidity(embeddingGeneratorKey, segments, nearestCacheContentVersionNode);
168        if (cacheContentVersionNode == null)
169        {
170            // Create new cache content version node
171            cacheContentVersionNode = _cacheDAO.getCacheContentVersionNode(content, revision);
172            
173            // Compare and replace segments under segments size
174            for (int i = 0; i < segments.size(); i++)
175            {
176                TextSegment textSegment = segments.get(i);
177                String newSegment = textSegment.text();
178                
179                CacheContentVersionChunk existingChunkNode = _cacheDAO.getChunk(cacheContentVersionNode, i, true);
180                String existingSegment = existingChunkNode.getValueOrDefault(DATA_CHUNK_CONTENT, null);
181                String existingCreator = existingChunkNode.getValueOrDefault(DATA_CHUNK_CREATOR, null);
182
183                CacheContentVersionChunk newChunkNode = _cacheDAO.getChunk(cacheContentVersionNode, i, true);
184
185                if (!Strings.CS.equals(newSegment, existingSegment)
186                        || !Strings.CS.equals(embeddingGeneratorKey, existingCreator))
187                {
188                    // compute and store embedding
189                    List<Float> embedding = _aiHelper.computeEmbedding(newSegment);
190                    if (embedding != null && !embedding.isEmpty())
191                    {
192                        newChunkNode.setValue(DATA_CHUNK_CONTENT, newSegment);
193                        newChunkNode.setValue(DATA_CHUNK_EMBDEDING, embedding.toArray(new Float[0]));
194                        newChunkNode.setValue(DATA_CHUNK_CREATION_DATETIME, ZonedDateTime.now());
195                        newChunkNode.setValue(DATA_CHUNK_CREATOR, embeddingGeneratorKey);
196                    }
197                }
198                else
199                {
200                    // simply copy old values
201                    newChunkNode.setValue(DATA_CHUNK_CONTENT, existingChunkNode.getValue(DATA_CHUNK_CONTENT));
202                    newChunkNode.setValue(DATA_CHUNK_EMBDEDING, existingChunkNode.getValue(DATA_CHUNK_EMBDEDING));
203                    newChunkNode.setValue(DATA_CHUNK_CREATION_DATETIME, existingChunkNode.getValue(DATA_CHUNK_CREATION_DATETIME));
204                    newChunkNode.setValue(DATA_CHUNK_CREATOR, existingChunkNode.getValue(DATA_CHUNK_CREATOR));
205                }
206            }
207            
208            if (cacheContentVersionNode.needsSave())
209            {
210                cacheContentVersionNode.saveChanges();
211            }
212        }
213        
214        return cacheContentVersionNode;
215    }
216
217    private List<SolrInputDocument> _getDocumentsFromCache(Content content, DefaultTraversableAmetysObject cacheContentVersionNode)
218    {
219        List<SolrInputDocument> sids = new ArrayList<>();
220        
221        // Read what is stored as vector and send it to solr
222        for (int i = 0; ; i++)
223        {
224            CacheContentVersionChunk chunk = _cacheDAO.getChunk(cacheContentVersionNode, i, false);
225            if (chunk == null)
226            {
227                break;
228            }
229            
230            Double[] v = chunk.getValueOrDefault(DATA_CHUNK_EMBDEDING, new Double[0]);
231            
232            SolrInputDocument doc = new SolrInputDocument();
233            doc.addField("id", chunk.getId());
234            doc.addField(SolrFieldNames.DOCUMENT_TYPE, SOLR_TYPE);
235            doc.addField(SOLR_EMBEDDING, Arrays.asList(v));
236            doc.addField(SOLR_CONTENT_ID, content.getId());
237            if (content instanceof WebContent webContent)
238            {
239                doc.addField(SolrWebFieldNames.SITE_NAME, webContent.getSiteName());
240            }
241            if (content.getLanguage() != null)
242            {
243                doc.addField(SolrFieldNames.CONTENT_LANGUAGES, content.getLanguage());
244            }
245            
246            _solrIndexer.indexAclInitValues(chunk, doc);
247            
248            sids.add(doc);
249        }
250
251        return sids;
252    }
253
254    private DefaultTraversableAmetysObject _checkCacheValidity(String embeddingGeneratorKey, List<TextSegment> segments, DefaultTraversableAmetysObject cacheContentVersionNode)
255    {
256        if (cacheContentVersionNode != null)
257        {
258            // Compare segments and segments size
259            for (int i = 0; i < segments.size(); i++)
260            {
261                TextSegment textSegment = segments.get(i);
262                String newSegment = textSegment.text();
263                
264                CacheContentVersionChunk chunkNode = _cacheDAO.getChunk(cacheContentVersionNode, i, false);
265                if (chunkNode == null)
266                {
267                    // Nearest has less segments, we will create a new one
268                    return null;
269                }
270                
271                String existingSegment = chunkNode.getValueOrDefault(DATA_CHUNK_CONTENT, null);
272                String creator = chunkNode.getValueOrDefault(DATA_CHUNK_CREATOR, null);
273                
274                if (!Strings.CS.equals(newSegment, existingSegment)
275                        || !Strings.CS.equals(embeddingGeneratorKey, creator))
276                {
277                    // Nearest is not matching, we will create a new one
278                    return null;
279                }
280            }
281        }
282        
283        // If cache is bigger that expected... that not a problem, we will ignore it
284        return cacheContentVersionNode;
285    }
286
287    private String _getContentRevision(Content content)
288    {
289        if (!(content instanceof VersionableAmetysObject versionableContent))
290        {
291            throw new IllegalArgumentException("Object should be a VersionableAmetysObject " + content.getId());
292        }
293            
294        String revision = versionableContent.getRevision();
295        if (revision != null)
296        {
297            return revision;
298        }
299        
300        Request request = ContextHelper.getRequest(_context);
301        String currentWsp = RequestAttributeWorkspaceSelector.getForcedWorkspace(request);
302        if (currentWsp.equals(WebConstants.LIVE_WORKSPACE))
303        {
304            String[] allLabels = versionableContent.getAllLabels();
305            if (Arrays.binarySearch(allLabels, "Live") >= 0)
306            {
307                versionableContent.switchToLabel("Live");
308                revision = versionableContent.getRevision();
309                versionableContent.switchToLabel(null);
310                return revision;
311            }
312            else
313            {
314                throw new IllegalArgumentException("Cannot determine live revision version of " + content.getId());
315            }
316        }
317        else
318        {
319            // get the last version
320            String[] allRevisions = versionableContent.getAllRevisions();
321            return allRevisions[allRevisions.length - 1];
322        }
323    }
324
325    public List<Query> getUnindexObjectQuery(String ametysObjectId) throws Exception
326    {
327        return List.of(new ContentChunkQuery(ametysObjectId));
328    }
329    
330    public Query getUnindexAllQuery(String type) throws Exception
331    {
332        return new DocumentTypeQuery(SOLR_TYPE);
333    }
334    
335    public Query getUnindexSiteQuery(String type, String sitename) throws Exception
336    {
337        return new AndQuery(getUnindexAllQuery(type), new SiteQuery(sitename));
338    }
339    
340    public Query getUnindexSitemapQuery(String type, String siteName, String sitemapName) throws Exception
341    {
342        return new AndQuery(getUnindexSiteQuery(type, siteName), new SitemapQuery(sitemapName));
343    }
344    
345    public Collection<SchemaDefinition> getSchemaDefinitions()
346    {
347        if (!_aiHelper.isEmbeddingSupported())
348        {
349            return List.of();
350        }
351        
352        org.apache.solr.client.solrj.request.schema.FieldTypeDefinition knnVectorFieldTypeDefinition = new org.apache.solr.client.solrj.request.schema.FieldTypeDefinition();
353        knnVectorFieldTypeDefinition.setAttributes(Map.of(
354            "name", "knn_vector",
355            "class", "solr.DenseVectorField",
356            "vectorDimension", _aiHelper.getEmbeddingDimension(),
357            "similarityFunction", "cosine"
358        ));
359        
360        return List.of(
361            new FieldTypeDefinition(knnVectorFieldTypeDefinition),
362            new FieldDefinition(SOLR_EMBEDDING, "knn_vector", false, false, true, false),
363            new FieldDefinition(SOLR_CONTENT_ID, "string", false, true, true, true)
364        );
365    }
366    
367    private List<TextSegment> _contentToSegments(Content content) throws IOException
368    {
369        String contentTextMarkdown = _contentToMD(content, "");
370        
371        TokenCountEstimator estimator = new DefaultTokenCountEstimator();
372        DocumentSplitter splitter = DocumentSplitters.recursive(__DOCUMENT_CHUNK_MAX_SIZE, __DOCUMENT_CHUNK_MAX_OVERLAP, estimator);
373        return StringUtils.isNotBlank(contentTextMarkdown) ? splitter.split(Document.from(contentTextMarkdown)) : List.of();
374    }
375    
376    private String _contentToMD (Content content, String version) throws IOException
377    {
378        if (!(content instanceof WebContent webContent))
379        {
380            return "";
381        }
382        
383        Request request = ContextHelper.getRequest(_context);
384        UserIdentity currentUser = AuthenticateAction.getUserIdentityFromSession(request);
385        Object currentInternalAllowed = request.getAttribute(AuthenticateAction.REQUEST_ATTRIBUTE_INTERNAL_ALLOWED);
386        
387        try
388        {
389            request.getSession(true).setAttribute(AuthenticateAction.SESSION_USERIDENTITY, UserPopulationDAO.SYSTEM_USER_IDENTITY);
390            request.setAttribute(AuthenticateAction.REQUEST_ATTRIBUTE_INTERNAL_ALLOWED, true);
391            
392            String uri = URIUtils.buildURI("cocoon://_wrapped-content-data.md", getContentViewUrlParameters(webContent, "ai", version));
393            SitemapSource src = null;
394            try
395            {
396                src = (SitemapSource) _sourceResolver.resolveURI(uri);
397                Reader reader = new InputStreamReader(src.getInputStream(), "UTF-8");
398                return IOUtils.toString(reader);
399            }
400            finally
401            {
402                _sourceResolver.release(src);
403            }
404        }
405        finally
406        {
407            request.setAttribute(AuthenticateAction.SESSION_USERIDENTITY, currentUser);
408            request.setAttribute(AuthenticateAction.REQUEST_ATTRIBUTE_INTERNAL_ALLOWED, currentInternalAllowed);
409        }
410    }
411    
412    private  Map<String, String> getContentViewUrlParameters(WebContent content, String viewName, String version)
413    {
414        Map<String, String> params = new HashMap<>();
415        params.put("contentId", content.getId());
416        params.put("viewName", viewName);
417        params.put("fallbackViewName", "main");
418        params.put("contentVersion", version);
419        params.put("showMissing", "false");
420        params.put("showDisableValues", "true");
421        
422        return params;
423    }
424}