Source code

001/*
002 *  Copyright 2025 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 */
016
017package org.ametys.plugins.ai.provider.impl;
018
019import java.io.IOException;
020import java.io.InputStream;
021import java.io.OutputStream;
022import java.net.URI;
023import java.nio.file.Files;
024import java.nio.file.Path;
025import java.util.ArrayList;
026import java.util.Arrays;
027import java.util.Base64;
028import java.util.List;
029
030import org.apache.avalon.framework.activity.Disposable;
031import org.apache.avalon.framework.activity.Initializable;
032import org.apache.avalon.framework.component.Component;
033import org.apache.avalon.framework.configuration.Configurable;
034import org.apache.avalon.framework.configuration.Configuration;
035import org.apache.avalon.framework.configuration.ConfigurationException;
036import org.apache.avalon.framework.service.ServiceException;
037import org.apache.avalon.framework.service.ServiceManager;
038import org.apache.commons.io.IOUtils;
039import org.apache.commons.lang3.StringUtils;
040import org.apache.hc.client5.http.classic.methods.HttpGet;
041import org.apache.hc.client5.http.impl.classic.CloseableHttpClient;
042
043import org.ametys.core.util.HttpUtils;
044import org.ametys.plugins.ai.AIHelper;
045import org.ametys.plugins.ai.provider.AIProvider;
046import org.ametys.runtime.config.Config;
047import org.ametys.runtime.i18n.I18nizableText;
048import org.ametys.runtime.plugin.component.AbstractLogEnabled;
049import org.ametys.runtime.plugin.component.DeferredServiceable;
050import org.ametys.runtime.plugin.component.PluginAware;
051import org.ametys.runtime.util.AmetysHomeHelper;
052
053import dev.langchain4j.data.document.Document;
054import dev.langchain4j.data.document.DocumentSplitter;
055import dev.langchain4j.data.document.splitter.DocumentSplitters;
056import dev.langchain4j.data.embedding.Embedding;
057import dev.langchain4j.data.image.Image;
058import dev.langchain4j.data.message.SystemMessage;
059import dev.langchain4j.data.message.UserMessage;
060import dev.langchain4j.data.segment.TextSegment;
061import dev.langchain4j.exception.AuthenticationException;
062import dev.langchain4j.exception.RateLimitException;
063import dev.langchain4j.memory.ChatMemory;
064import dev.langchain4j.model.TokenCountEstimator;
065import dev.langchain4j.model.chat.ChatModel;
066import dev.langchain4j.model.chat.request.ChatRequest;
067import dev.langchain4j.model.chat.response.ChatResponse;
068import dev.langchain4j.model.embedding.EmbeddingModel;
069import dev.langchain4j.model.image.ImageModel;
070import dev.langchain4j.model.output.Response;
071
072/**
073 * This is an abstract class for AI providers.
074 */
075public abstract class AbstractAIProvider extends AbstractLogEnabled implements AIProvider, Configurable, Initializable, PluginAware, DeferredServiceable, Component, Disposable
076{
077    /** The configuration parameter for the API key  */
078    public static final String CONFIG_API_KEY = "ai.apikey";
079    
080    /** The configuration parameter for URL */
081    public static final String CONFIG_SERVER_URL = "ai.serverURL";
082    
083    /** The configuration parameter for text model */
084    public static final String CONFIG_TEXT_MODEL = "ai.model.text";
085    
086    /** The configuration parameter for image model */
087    public static final String CONFIG_IMAGE_MODEL = "ai.model.image";
088    
089    /** The configuration parameter for image generation timeout */
090    public static final String CONFIG_IMAGE_TIMEOUT = "ai.model.image.timeout";
091
092    /** The configuration parameter for embedding model */
093    public static final String CONFIG_EMBEDDING_MODEL = "ai.model.embedding";
094    
095    /** The provider id */
096    protected String _id;
097    /** The provider label */
098    protected I18nizableText _label;
099    /** The plugin name */
100    protected String _pluginName;
101    
102    /** The provider models existing for text generation */
103    protected List<String> _textModelNames;
104    /** The provider models existing for image generation */
105    protected List<String> _imageModelNames;
106    /** The provider models existing for embedding */
107    protected List<EmbeddingModelInfo> _embeddingModelInfos;
108    
109    /** Current text model */
110    protected ChatModel _textModel;
111    /** The context window for text generation */
112    protected int _textModelContextWindow;
113    /** Current image model */
114    protected ImageModel _imageModel;
115    /** The context window for image generation */
116    protected int _imageModelContextWindow;
117    /** Current embedding model */
118    protected EmbeddingModel _embeddingModel;
119    
120    /** The AI helper */
121    protected AIHelper _aiHelper;
122    
123    /** The HTTP client used to fetch images */
124    protected CloseableHttpClient _httpClient;
125
126    /**
127     * Info for an embedding model
128     * @param name The name of the model
129     * @param vectorDimension The vector dimension
130     */
131    public record EmbeddingModelInfo(String name, int vectorDimension) { }
132    
133
134    public void setPluginInfo(String pluginName, String featureName, String id)
135    {
136        _pluginName = pluginName;
137        _id = id;
138    }
139
140    public void configure(Configuration configuration) throws ConfigurationException
141    {
142        _label = I18nizableText.parseI18nizableText(configuration.getChild("label"), "plugin." + _pluginName, "");
143        
144        Configuration modelConfiguration = configuration.getChild("model");
145        
146        // Text
147        Configuration textConfiguration = modelConfiguration.getChild("text");
148        
149        _textModelNames = Arrays.stream(textConfiguration.getChildren("name"))
150            .map(c -> c.getValue(null))
151            .filter(StringUtils::isNotBlank)
152            .toList();
153        
154        int contextWindow = textConfiguration.getChild("contextWindow").getValueAsInteger(-1);
155        if (contextWindow == -1)
156        {
157            throw new ConfigurationException("Missing <contextWindow> for model <text>", modelConfiguration);
158        }
159        _textModelContextWindow = contextWindow;
160
161        // Image
162        Configuration imageConfiguration = modelConfiguration.getChild("image", false);
163        if (imageConfiguration != null)
164        {
165            _imageModelNames = Arrays.stream(imageConfiguration.getChildren("name"))
166                    .map(c -> c.getValue(null))
167                    .filter(StringUtils::isNotBlank)
168                    .toList();
169                
170            contextWindow = imageConfiguration.getChild("contextWindow").getValueAsInteger(-1);
171            if (contextWindow == -1)
172            {
173                throw new ConfigurationException("Missing <contextWindow> for model <image>", modelConfiguration);
174            }
175            
176            _imageModelContextWindow = contextWindow;
177        }
178        else
179        {
180            _imageModelNames = List.of();
181        }
182
183
184        // Embedding
185        Configuration embeddingConfiguration = modelConfiguration.getChild("embedding", false);
186        if (embeddingConfiguration != null)
187        {
188            _embeddingModelInfos = Arrays.stream(embeddingConfiguration.getChildren("name"))
189                    .map(c -> new EmbeddingModelInfo(c.getValue(null), Integer.parseInt(c.getAttribute("vectorDimension", ""))))
190                    .filter(v -> StringUtils.isNotBlank(v.name()))
191                    .toList();
192        }
193        else
194        {
195            _embeddingModelInfos = List.of();
196        }
197    }
198    
199    private synchronized void _initModels()
200    {
201        if (_textModel == null)
202        {
203            _textModel = createChatModel(Config.getInstance().getValue(CONFIG_TEXT_MODEL));
204            
205            if (_imageModelContextWindow != -1)
206            {
207                String imageModelName = Config.getInstance().getValue(CONFIG_IMAGE_MODEL);
208                if (StringUtils.isNotBlank(imageModelName))
209                {
210                    _imageModel = createImageModel(imageModelName);
211                }
212            }
213
214            if (_embeddingModelInfos.size() > 0)
215            {
216                String embeddingModelName = Config.getInstance().getValue(CONFIG_EMBEDDING_MODEL);
217                if (StringUtils.isNotBlank(embeddingModelName))
218                {
219                    _embeddingModel = createEmbeddingModel(embeddingModelName);
220                }
221            }
222        }
223    }
224    
225    /**
226     * Creates the text model implementation
227     * @param modelName The model name
228     * @return The chat model
229     */
230    protected abstract ChatModel createChatModel(String modelName);
231    
232    /**
233     * Creates the image model implementation
234     * @param modelName The model name
235     * @return The image model
236     */
237    protected ImageModel createImageModel(String modelName)
238    {
239        throw new UnsupportedOperationException(this.getClass().getName() + " does not support image generation");
240    }
241    
242    /**
243     * Creates the embedding model implementation
244     * @param modelName The model name
245     * @return The embedding model
246     */
247    protected abstract EmbeddingModel createEmbeddingModel(String modelName);
248
249    public void initialize() throws Exception
250    {
251        _httpClient = HttpUtils.createHttpClient(0, 2, false);
252    }
253    
254    public void dispose()
255    {
256        try
257        {
258            _httpClient.close();
259        }
260        catch (IOException e)
261        {
262            throw new RuntimeException(e);
263        }
264    }
265    
266    public void deferredService(ServiceManager manager) throws ServiceException
267    {
268        _aiHelper = (AIHelper) manager.lookup(AIHelper.ROLE);
269    }
270    
271    public String getId()
272    {
273        return _id;
274    }
275
276    public I18nizableText getLabel()
277    {
278        return _label;
279    }
280    
281    /**
282     * Get the list of known text models for this provider
283     * @return The list
284     */
285    public List<String> getKnownTextModels()
286    {
287        return _textModelNames;
288    }
289
290    /**
291     * Get the list of known image models for this provider
292     * @return The non null list
293     */
294    public List<String> getKnownImageModels()
295    {
296        return _imageModelNames;
297    }
298    
299    /**
300     * Get the list of known embedding models for this provider
301     * @return the non null list
302     */
303    public List<EmbeddingModelInfo> getKnownEmbeddingModels()
304    {
305        return _embeddingModelInfos;
306    }
307    
308    /**
309     * Creates a token count estimator
310     * @return The non null token count estimator
311     */
312    protected TokenCountEstimator _getTokenCountEstimator()
313    {
314        return new DefaultTokenCountEstimator();
315    }
316
317    public String textToSummary(String text, int maxLength) throws Exception
318    {
319        _initModels();
320        
321        TokenCountEstimator estimator = _getTokenCountEstimator();
322        
323        // Get the token number of the prompt without the input text, to not exceed context window.
324        int tokenForPrompt = estimator.estimateTokenCountInText(_aiHelper.getTextSummaryPrompt(maxLength));
325        int contextWindow = _textModelContextWindow;
326        int tokenMax = contextWindow - tokenForPrompt;
327        
328        List<String> summaries = new ArrayList<>();
329        DocumentSplitter splitter = DocumentSplitters.recursive(tokenMax , 5, estimator);
330        List<TextSegment> segments = splitter.split(Document.from(text));
331        for (TextSegment segment : segments)
332        {
333            boolean success = false;
334            int retries = 0;
335            while (!success && retries < 3)
336            {
337                try
338                {
339                    ChatResponse response = _textModel.chat(new SystemMessage(_aiHelper.getTextSummaryPrompt(maxLength)), new UserMessage(segment.text()));
340                    summaries.add(response.aiMessage().text());
341                    success = true;
342                }
343                catch (RateLimitException e)
344                {
345                    retries++;
346                    Thread.sleep(40000); // Waiting a little to low the token per minute
347                }
348            }
349            
350            if (!success)
351            {
352                getLogger().warn("An error occurred during summary generation of the segment. So we ignore it.");
353            }
354        }
355        
356        // Only if one summary exists, return the summary
357        if (summaries.size() == 1)
358        {
359            return summaries.get(0);
360        }
361        // Else, return a final summary of all summaries
362        String summary = String.join(" ", summaries);
363        String prompt = _aiHelper.getTextSummaryPrompt(maxLength);
364        
365        ChatResponse response = _textModel.chat(SystemMessage.from(prompt), UserMessage.from(summary));
366        return response.aiMessage().text();
367    }
368    
369    public boolean isImageGenerationSupported()
370    {
371        _initModels();
372        
373        return _imageModel != null;
374    }
375    
376    public Path textToImage(String text) throws Exception
377    {
378        _initModels();
379        
380        TokenCountEstimator estimator = new DefaultTokenCountEstimator();
381        
382        // Get the token number of the prompt without the input text, to not exceed context window.
383        int tokenForPrompt = estimator.estimateTokenCountInText(_aiHelper.getImageGenerationPrompt());
384        int contextWindow = _imageModelContextWindow;
385        int tokenMax = contextWindow - tokenForPrompt;
386        
387        DocumentSplitter splitter = DocumentSplitters.recursive(tokenMax, 0, estimator);
388        List<TextSegment> segments = splitter.split(Document.from(text));
389
390        if (segments.isEmpty())
391        {
392            throw new IllegalArgumentException("An error occurred generating the image. Can't build a segment for the given text " + text);
393        }
394        String res = segments.get(0).text();
395
396        String prompt = _aiHelper.getImageGenerationPrompt();
397        
398        Response<Image> response = _imageModel.generate(prompt + " Here is the user text: " + res);
399        Image image = response.content();
400        
401        if (image.base64Data() != null)
402        {
403            String tmpFolderName = Long.toString(Math.round(Math.random() * 1000000.0));
404            String fileName = "image.png";
405            
406            Path targetFile = AmetysHomeHelper.getAmetysHomeTmp().toPath().resolve(tmpFolderName).resolve(fileName);
407            Files.createDirectories(targetFile.getParent());
408
409            byte[] imageBytes = Base64.getDecoder().decode(image.base64Data());
410            Files.write(targetFile, imageBytes);
411            return targetFile;
412        }
413        else
414        {
415            return _download(image.url());
416        }
417    }
418    
419    /**
420     * Helper to download the given url to a temporary file
421     * @param uri The uri to download
422     * @return The created file
423     * @throws IOException If an error occurred
424     */
425    protected Path _download(URI uri) throws IOException
426    {
427        HttpGet httpGet = new HttpGet(uri);
428        
429        String tmpFolderName = Long.toString(Math.round(Math.random() * 1000000.0));
430        String fileName = StringUtils.substringAfterLast(uri.getPath(), "/");
431        
432        Path targetFile = AmetysHomeHelper.getAmetysHomeTmp().toPath().resolve(tmpFolderName).resolve(fileName);
433        Files.createDirectories(targetFile.getParent());
434        
435        return _httpClient.execute(httpGet, httpResponse ->
436        {
437            if (httpResponse.getCode() != 200)
438            {
439                throw new IOException("Error " + httpResponse.getCode() + " while accessing to the URI " + uri);
440            }
441            
442            try (InputStream is = httpResponse.getEntity().getContent();
443                    OutputStream os = Files.newOutputStream(targetFile))
444            {
445                IOUtils.copy(is, os);
446            }
447            
448            return targetFile;
449        });
450    }
451
452    
453    /**
454     * Get the API key from the configuration
455     * @return the API key
456     */
457    protected String getAPIKey()
458    {
459        return Config.getInstance().getValue(CONFIG_API_KEY);
460    }
461
462    /**
463     * Get the server URL from the configuration
464     * @return the server URL
465     */
466    protected String getServerURL()
467    {
468        return Config.getInstance().getValue(CONFIG_SERVER_URL);
469    }
470    
471    public boolean isEmbeddingSupported()
472    {
473        _initModels();
474        
475        return _embeddingModel != null;
476    }
477    
478    public String askChatbot(ChatMemory memory)
479    {
480        _initModels();
481        
482        ChatRequest chatRequest = ChatRequest.builder()
483                .messages(memory.messages())
484                .build();
485        
486        ChatResponse chatResponse = _textModel.chat(chatRequest);
487        return chatResponse.aiMessage().text();
488    }
489    
490    public List<Float> computeEmbedding(String text)
491    {
492        try
493        {
494            _initModels();
495            
496            Embedding embedding = _embeddingModel.embed(text).content();
497            List<Float> embeddingList = embedding.vectorAsList();
498            
499            if (isEmbeddingNormalizationRequired())
500            {
501                return normalize(embeddingList);
502            }
503            else
504            {
505                return embeddingList;
506            }
507        }
508        catch (AuthenticationException e)
509        {
510            getLogger().warn("Authentication error while computing embedding. Please check the API key.", e);
511            return null;
512        }
513    }
514    
515    /**
516     * Is normalization required for the embeddings of this provider
517     * @return true if normalization is required
518     */
519    protected boolean isEmbeddingNormalizationRequired()
520    {
521        return false;
522    }
523
524    /**
525     * Necessary for open models who often are not normalized.
526     * @param vector the vector to normalize
527     * @return the normalized vector
528     */
529    public List<Float> normalize (List<Float> vector)
530    {
531        double sumSquares = 0.0;
532        for (Float v : vector)
533        {
534            sumSquares += v * v;
535        }
536        double norm = Math.sqrt(sumSquares);
537        if (norm == 0)
538        {
539            return vector;
540        }
541        List<Float> normalized = new ArrayList<>(vector.size());
542        for (Float v : vector)
543        {
544            normalized.add((float) (v / norm));
545        }
546        return normalized;
547    }
548    
549    public int getEmbeddingDimension()
550    {
551        String modelName = Config.getInstance().getValue(CONFIG_EMBEDDING_MODEL);
552        if (StringUtils.isNotBlank(modelName))
553        {
554            for (EmbeddingModelInfo info : _embeddingModelInfos)
555            {
556                if (info.name().equals(modelName))
557                {
558                    return info.vectorDimension();
559                }
560            }
561        }
562        return -1;
563    }
564    
565    public String getEmbeddingKey()
566    {
567        if (isEmbeddingSupported())
568        {
569            return this.getId() + "_" + Config.getInstance().getValue(CONFIG_EMBEDDING_MODEL) + ":"  + getEmbeddingDimension();
570        }
571        return null;
572    }
573}