from global_vars import translations, t from app import Plugin import streamlit as st import yaml from litellm import completion, embedding import numpy as np from sklearn.metrics.pairwise import cosine_similarity, euclidean_distances, manhattan_distances import os from typing import List, Dict, Any import requests import torch from transformers import AutoTokenizer, AutoModel from huggingface_hub import InferenceClient from langchain_huggingface import HuggingFaceEmbeddings DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' MAX_LENGTH = 512 CHUNK_SIZE = 200 def mean_pooling(model_output, attention_mask): token_embeddings = model_output[0] input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9) def mean_pooling(model_output, attention_mask): token_embeddings = model_output[0] input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9) # Ajout des traductions spécifiques à ce plugin translations["en"].update({ "rag_plugin_loaded": "RAG LLM Plugin loaded", "rag_enter_text": "Enter RAG text:", "rag_enter_question": "Enter your question:", "rag_button_get_answer": "Get an answer", "rag_success_text_processed": "RAG text processed successfully!", "rag_warning_enter_text": "Please enter RAG text.", "rag_warning_process_text_first": "Please process the RAG text first.", "rag_warning_enter_question": "Please enter a question.", "rag_answer": "Answer:", "rag_citations": "Citations:", "rag_model_provider": "Model Provider", "rag_llm_model": "LLM Model", "rag_embedder_model": "Embedding Model", "rag_similarity_method": "Similarity Method", "rag_llm_sys_prompt": "System prompt for LLM", "rag_chunk_size": "Chunk size", "rag_top_k_chunks": "Number of chunks to use", "rag_default_sys_prompt": "You are an AI assistant. Your task is to analyze the provided context and answer questions based ONLY on this context. If the information is not in the context, clearly state that.", "rag_error_fetching_models_ollama": "Error fetching Ollama models: ", "rag_error_calling_llm": "Error calling LLM: ", "rag_processing" : "Processing...", "rag_hf_api_key": "HuggingFace API Token", }) translations["fr"].update({ "rag_plugin_loaded": "Plugin RAG LLM chargé", "rag_enter_text": "Entrez le texte RAG :", "rag_enter_question": "Entrez votre question :", "rag_button_get_answer": "Obtenir une réponse", "rag_success_text_processed": "Texte RAG traité avec succès!", "rag_warning_enter_text": "Veuillez entrer du texte RAG.", "rag_warning_process_text_first": "Veuillez d'abord traiter le texte RAG.", "rag_warning_enter_question": "Veuillez entrer une question.", "rag_answer": "Réponse :", "rag_citations": "Citations :", "rag_model_provider": "Fournisseur de modèle", "rag_llm_model": "Modèle LLM", "rag_embedder_model": "Modèle d'embedding", "rag_similarity_method": "Méthode de similarité", "rag_llm_sys_prompt": "Prompt système pour le LLM", "rag_chunk_size": "Taille des chunks", "rag_top_k_chunks": "Nombre de chunks à utiliser", "rag_default_sys_prompt": "Tu es un assistant IA. Ta tâche est d'analyser le contexte fourni et de répondre aux questions en te basant UNIQUEMENT sur ce contexte. Si l'information n'est pas dans le contexte, dis-le clairement.", "rag_error_fetching_models_ollama": "Erreur lors de la récupération des modèles Ollama : ", "rag_error_calling_llm": "Erreur lors de l'appel au LLM : ", "rag_processing" : "En cours de traitement...", "rag_hf_api_key": "Token API HuggingFace", }) class RagllmPlugin(Plugin): def __init__(self, name: str, plugin_manager): super().__init__(name, plugin_manager) try: self.config = self.load_llm_config() except: self.config = {} self.embeddings = None self.chunks = None self.hf_client = None def load_llm_config(self) -> Dict: try: with open('.llm-config.yml', 'r') as file: return yaml.safe_load(file) except: return {} def get_tabs(self): return [{"name": "RAG", "plugin": "ragllm"}] def get_config_fields(self): fields = { "provider": { "type": "select", "label": t("rag_model_provider"), "options": [("ollama", "Ollama"), ("groq", "Groq"), ("huggingface", "HuggingFace")], "default": "ollama" }, "llm_model": { "type": "select", "label": t("rag_llm_model"), "options": [("none", "À charger...")], "default": "ollama/qwen2" }, "embedder": { "type": "select", "label": t("rag_embedder_model"), "options": [ ("sentence-transformers/all-MiniLM-L6-v2", "all-MiniLM-L6-v2"), ("nomic-ai/nomic-embed-text-v1.5", "nomic-embed-text-v1.5") ], "default": "sentence-transformers/all-MiniLM-L6-v2" }, "similarity_method": { "type": "select", "label": t("rag_similarity_method"), "options": [ ("cosine", "Cosinus"), ("euclidean", "Distance euclidienne"), ("manhattan", "Distance de Manhattan") ], "default": "cosine" }, "llm_sys_prompt": { "type": "textarea", "label": t("rag_llm_sys_prompt"), "default": t("rag_default_sys_prompt") }, "chunk_size": { "type": "number", "label": t("rag_chunk_size"), "default": 200 }, "top_k": { "type": "number", "label": t("rag_top_k_chunks"), "default": 3 } } # Add HuggingFace API key field if provider is huggingface if 'provider' in self.config and self.config.get('provider') == 'huggingface': fields["hf_api_key"] = { "type": "password", "label": t("rag_hf_api_key"), "default": "" } return fields def get_config_ui(self, config): updated_config = {} for field, params in self.get_config_fields().items(): if params['type'] == 'select': if field == 'llm_model': provider = config.get('provider', 'ollama') models = self.get_available_models(provider) try: default_index = models.index(config.get(field, params['default'])) except ValueError: default_index = 0 updated_config[field] = st.selectbox( params['label'], options=models, index=default_index ) else: options_list = [option[0] for option in params['options']] try: default_index = options_list.index(config.get(field, params['default'])) except ValueError: default_index = 0 updated_config[field] = st.selectbox( params['label'], options=options_list, format_func=lambda x: dict(params['options'])[x], index=default_index ) elif params['type'] == 'textarea': updated_config[field] = st.text_area( params['label'], value=config.get(field, params['default']) ) elif params['type'] == 'number': updated_config[field] = st.number_input( params['label'], value=int(config.get(field, params['default'])), step=1 ) else: updated_config[field] = st.text_input( params['label'], value=config.get(field, params['default']) ) return updated_config def get_sidebar_config_ui(self, config: Dict[str, Any]) -> Dict[str, Any]: available_models = self.get_available_models('ollama') + self.get_available_models('groq') default_model = config.get('llm_model', available_models[0] if available_models else None) selected_model = st.sidebar.selectbox( t("rag_llm_model"), options=available_models, index=available_models.index(default_model) if default_model in available_models else 0, key="ragllm_llm_model" ) return {"llm_model": selected_model} def get_available_models(self, provider: str) -> List[str]: if provider == 'ollama': try: response = requests.get("http://localhost:11434/api/tags") models = response.json()["models"] return [f"ollama/{model['name']}" for model in models] + ["ollama/qwen2"] except Exception as e: st.error(f"{t('rag_error_fetching_models_ollama')}{str(e)}") return ["ollama/qwen2"] elif provider == 'groq': return ["groq/llama3-70b-8192", "groq/mixtral-8x7b-32768"] elif provider == 'huggingface': return ["HuggingFaceH4/zephyr-7b-beta"] else: return ["none"] def process_rag_text(self, rag_text: str, chunk_size: int, embedder): rag_text = rag_text.replace('\\n', ' ').replace('\\\'', "'") mots = rag_text.split() self.chunks = [' '.join(mots[i:i+chunk_size]) for i in range(0, len(mots), chunk_size)] self.embeddings = np.vstack([self.get_embedding(c, embedder) for c in self.chunks]) def get_embedding(self, text: str, model: str) -> np.ndarray: if self.config.get('provider') == 'huggingface': if not hasattr(self, 'hf_embeddings'): self.hf_embeddings = HuggingFaceEmbeddings( model_name=model, task="feature-extraction", encode_kwargs={'normalize': True} ) embedding = self.hf_embeddings.embed_query(text) return np.array(embedding).reshape(1, -1) else: # Original embedding logic tokenizer = AutoTokenizer.from_pretrained(model) model = AutoModel.from_pretrained(model, trust_remote_code=True).to(DEVICE) inputs = tokenizer(text, padding=True, truncation=True, max_length=MAX_LENGTH, return_tensors="pt").to(DEVICE) with torch.no_grad(): model_output = model(**inputs) return mean_pooling(model_output, inputs['attention_mask']).cpu().numpy() def calculate_similarity(self, query_embedding: np.ndarray, method: str) -> np.ndarray: if method == 'cosine': return cosine_similarity(query_embedding.reshape(1, -1), self.embeddings)[0] elif method == 'euclidean': return -euclidean_distances(query_embedding.reshape(1, -1), self.embeddings)[0] elif method == 'manhattan': return -manhattan_distances(query_embedding.reshape(1, -1), self.embeddings)[0] else: raise ValueError("Méthode de similarité non reconnue") def get_context(self, query: str, config: Dict[str, Any]) -> tuple: query_embedding = self.get_embedding(query, config['ragllm']['embedder']) similarities = self.calculate_similarity(query_embedding, config['ragllm']['similarity_method']) top_indices = np.argsort(similarities)[-config['ragllm']['top_k']:][::-1] context = "\n\n".join([self.chunks[i] for i in top_indices]) return context, [self.chunks[i] for i in top_indices] def call_llm(self, prompt: str, sysprompt: str) -> str: try: llm_model = st.session_state.ragllm_llm_model if self.config.get('provider') == 'huggingface': if not self.hf_client: self.hf_client = InferenceClient(token=self.config.get('hf_api_key')) messages = [ {"role": "system", "content": sysprompt}, {"role": "user", "content": prompt} ] response = self.hf_client.text_generation( model=llm_model, prompt=prompt, max_new_tokens=512, temperature=0.7, stream=False ) return response else: messages = [ {"role": "system", "content": sysprompt}, {"role": "user", "content": prompt} ] response = completion(model=llm_model, messages=messages) return response['choices'][0]['message']['content'] except Exception as e: return f"{t('rag_error_calling_llm')}{str(e)}" def free_llm(self): try: llm_model = st.session_state.ragllm_llm_model if llm_model.startswith("ollama/"): ollama_model = llm_model.split("/")[1] response = requests.post( "http://localhost:11434/api/generate", json={ "model": ollama_model, "prompt": "bye", "keep_alive": 0 } ) return response.json()['response'] except Exception as e: return f"{t('rag_error_calling_llm')}{str(e)}" def process_with_llm(self, prompt: str, sysprompt: str, context: str) -> str: return self.call_llm(f"Contexte : {context}\n\nQuestion : {prompt}", sysprompt) def run(self, config): st.write(t("rag_plugin_loaded")) # Initialiser rag_text avec la valeur de session_state si elle existe, sinon utiliser une chaîne vide if 'rag_text' not in st.session_state: st.session_state.rag_text = "" if 'rag_question' not in st.session_state: st.session_state.rag_question = "Question" rag_text = st.text_area(t("rag_enter_text"), height=200, value=st.session_state.rag_text, key="rag_text_key") user_prompt = st.text_area(t("rag_enter_question"), value=st.session_state.rag_question, key="rag_prompt_key") st.session_state.rag_text = rag_text # Mettre à jour la valeur dans session_state st.session_state.rag_question = user_prompt if st.button(t("rag_button_get_answer"), key="get_answer_button"): with st.spinner(t("rag_processing")): if rag_text: self.process_rag_text(rag_text, config['ragllm']['chunk_size'], config['ragllm']['embedder']) st.success(t("rag_success_text_processed")) else: st.warning(t("rag_warning_enter_text")) if user_prompt and self.embeddings is not None: context, citations = self.get_context(user_prompt, config) response = self.process_with_llm(user_prompt, config['ragllm']['llm_sys_prompt'], context) st.write(t("rag_answer")) st.write(response) st.write(t("rag_citations")) for i, citation in enumerate(citations, 1): st.write(f"{i}. {citation[:100]}...") elif self.embeddings is None: st.warning(t("rag_warning_process_text_first")) else: st.warning(t("rag_warning_enter_question"))