413 lines
17 KiB
Python
413 lines
17 KiB
Python
import traceback
|
|
import threading
|
|
import queue
|
|
from langchain.prompts import ChatPromptTemplate
|
|
from langchain_ollama import ChatOllama
|
|
from rag_chatbot import MultimodalRAGChatbot
|
|
from config.settings import QDRANT_URL, QDRANT_COLLECTION_NAME, EMBEDDING_MODEL, OLLAMA_URL, DEFAULT_MODEL
|
|
from translations.lang_mappings import LANGUAGE_MAPPING
|
|
from utils.image_utils import base64_to_image
|
|
from langchain.callbacks.base import BaseCallbackHandler
|
|
import re
|
|
|
|
def clean_llm_response(text):
|
|
"""Nettoie la réponse du LLM en enlevant les balises de pensée et autres éléments non désirés."""
|
|
# Supprimer les blocs de pensée (<think>...</think>)
|
|
text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
|
|
# Supprimer les espaces supplémentaires au début de la réponse
|
|
text = text.lstrip()
|
|
return text
|
|
# Handler personnalisé pour le streaming
|
|
class GradioStreamingHandler(BaseCallbackHandler):
|
|
def __init__(self):
|
|
self.tokens_queue = queue.Queue()
|
|
self.full_text = ""
|
|
|
|
def on_llm_new_token(self, token, **kwargs):
|
|
self.tokens_queue.put(token)
|
|
self.full_text += token
|
|
|
|
# Initialiser le chatbot
|
|
rag_bot = MultimodalRAGChatbot(
|
|
qdrant_url=QDRANT_URL,
|
|
qdrant_collection_name=QDRANT_COLLECTION_NAME,
|
|
ollama_model=DEFAULT_MODEL,
|
|
embedding_model=EMBEDDING_MODEL,
|
|
ollama_url=OLLAMA_URL
|
|
)
|
|
print(f"Chatbot initialisé avec modèle: {DEFAULT_MODEL}")
|
|
|
|
# Variables globales
|
|
current_images = []
|
|
current_tables = []
|
|
|
|
# Fonctions utilitaires
|
|
def display_images(images_list=None):
|
|
"""Crée une liste de tuples (image, caption) pour Gradio Gallery"""
|
|
images_to_use = images_list if images_list is not None else current_images
|
|
|
|
if not images_to_use:
|
|
return None
|
|
|
|
gallery = []
|
|
for img_data in images_to_use:
|
|
image = img_data["image"]
|
|
if image:
|
|
caption = f"{img_data['caption']} (Source: {img_data['source']}, Page: {img_data['page']})"
|
|
gallery.append((image, caption))
|
|
|
|
return gallery if gallery else None
|
|
|
|
def display_tables(tables_list=None, language=None):
|
|
"""Crée le HTML pour afficher les tableaux"""
|
|
tables_to_use = tables_list if tables_list is not None else current_tables
|
|
|
|
if not tables_to_use:
|
|
return None
|
|
|
|
html = ""
|
|
for idx, table in enumerate(tables_to_use):
|
|
table_data = table['data']
|
|
table_html = ""
|
|
|
|
try:
|
|
if isinstance(table_data, str):
|
|
if '|' in table_data:
|
|
rows = table_data.strip().split('\n')
|
|
table_html = '<div class="table-container"><table>'
|
|
|
|
for i, row in enumerate(rows):
|
|
if i == 1 and all(c in ':-|' for c in row):
|
|
continue
|
|
|
|
cells = row.split('|')
|
|
|
|
if cells and cells[0].strip() == '':
|
|
cells = cells[1:]
|
|
if cells and cells[-1].strip() == '':
|
|
cells = cells[:-1]
|
|
|
|
if cells:
|
|
is_header = (i == 0)
|
|
table_html += '<tr>'
|
|
for cell in cells:
|
|
cell_content = cell.strip()
|
|
if is_header:
|
|
table_html += f'<th>{cell_content}</th>'
|
|
else:
|
|
table_html += f'<td>{cell_content}</td>'
|
|
table_html += '</tr>'
|
|
|
|
table_html += '</table></div>'
|
|
else:
|
|
table_html = f'<pre>{table_data}</pre>'
|
|
else:
|
|
table_html = f'<pre>{table_data}</pre>'
|
|
except Exception as e:
|
|
print(f"Error formatting table {idx}: {e}")
|
|
table_html = f'<pre>{table_data}</pre>'
|
|
|
|
html += f"""
|
|
<div style="margin-bottom: 20px; border: 1px solid #ddd; padding: 15px; border-radius: 8px;">
|
|
<h3>{table.get('caption', 'Tableau')}</h3>
|
|
<p style="color:#666; font-size:0.9em;">Source: {table.get('source', 'N/A')}, Page: {table.get('page', 'N/A')}</p>
|
|
<p><strong>Description:</strong> {table.get('description', '')}</p>
|
|
{table_html}
|
|
</div>
|
|
"""
|
|
|
|
return html if html else None
|
|
|
|
# Fonction pour changer de modèle
|
|
def change_model(model_name, language="Français"):
|
|
global rag_bot
|
|
|
|
try:
|
|
rag_bot = MultimodalRAGChatbot(
|
|
qdrant_url=QDRANT_URL,
|
|
qdrant_collection_name=QDRANT_COLLECTION_NAME,
|
|
ollama_model=model_name,
|
|
embedding_model=EMBEDDING_MODEL,
|
|
ollama_url=OLLAMA_URL
|
|
)
|
|
print(f"Modèle changé pour: {model_name}")
|
|
return f"✅ Modèle changé pour: {model_name}"
|
|
except Exception as e:
|
|
print(f"Erreur lors du changement de modèle: {e}")
|
|
return f"❌ Erreur: {str(e)}"
|
|
|
|
# Fonction pour changer de collection
|
|
def change_collection(collection_name, language="Français"):
|
|
global rag_bot
|
|
|
|
try:
|
|
rag_bot = MultimodalRAGChatbot(
|
|
qdrant_url=QDRANT_URL,
|
|
qdrant_collection_name=collection_name,
|
|
ollama_model=rag_bot.llm.model,
|
|
embedding_model=EMBEDDING_MODEL,
|
|
ollama_url=OLLAMA_URL
|
|
)
|
|
print(f"Collection changée pour: {collection_name}")
|
|
return f"✅ Collection changée pour: {collection_name}"
|
|
except Exception as e:
|
|
print(f"Erreur lors du changement de collection: {e}")
|
|
return f"❌ Erreur: {str(e)}"
|
|
|
|
# Fonction de traitement de requête
|
|
def convert_to_messages_format(history):
|
|
"""Convertit différents formats d'historique au format messages."""
|
|
messages = []
|
|
|
|
# Vérifier si nous avons déjà le format messages
|
|
if history and isinstance(history[0], dict) and "role" in history[0]:
|
|
return history
|
|
|
|
# Format tuples [(user_msg, assistant_msg), ...]
|
|
try:
|
|
for item in history:
|
|
if isinstance(item, tuple) and len(item) == 2:
|
|
user_msg, assistant_msg = item
|
|
messages.append({"role": "user", "content": user_msg})
|
|
if assistant_msg: # Éviter les messages vides
|
|
messages.append({"role": "assistant", "content": assistant_msg})
|
|
except ValueError:
|
|
# Journaliser l'erreur pour le débogage
|
|
print(f"Format d'historique non reconnu: {history}")
|
|
# Retourner un historique vide en cas d'erreur
|
|
return []
|
|
|
|
return messages
|
|
|
|
def process_query(message, history, streaming, show_sources, max_images, language):
|
|
global current_images, current_tables
|
|
|
|
# Debug plus clair
|
|
print(f"Langue sélectionnée pour la réponse: {language} -> {LANGUAGE_MAPPING.get(language, 'français')}")
|
|
|
|
if not message.strip():
|
|
return history, "", None, None
|
|
|
|
current_images = []
|
|
current_tables = []
|
|
print(f"Traitement du message: {message}")
|
|
print(f"Streaming: {streaming}")
|
|
|
|
try:
|
|
if streaming:
|
|
# Convertir history en format messages pour l'affichage
|
|
messages_history = convert_to_messages_format(history)
|
|
messages_history.append({"role": "user", "content": message})
|
|
messages_history.append({"role": "assistant", "content": ""})
|
|
|
|
# 1. Récupérer les documents pertinents
|
|
docs = rag_bot._retrieve_relevant_documents(message)
|
|
|
|
# 2. Préparer le contexte et l'historique
|
|
context = rag_bot._format_documents(docs)
|
|
history_text = rag_bot._format_chat_history()
|
|
|
|
# 3. Préparer le prompt
|
|
prompt_template = ChatPromptTemplate.from_template("""
|
|
Tu es un assistant documentaire spécialisé qui utilise le contexte fourni.
|
|
|
|
===== INSTRUCTION CRUCIALE SUR LA LANGUE =====
|
|
RÉPONDS UNIQUEMENT EN {language}. C'est une exigence ABSOLUE.
|
|
NE RÉPONDS JAMAIS dans une autre langue que {language}, quelle que soit la langue de la question.
|
|
==============================================
|
|
|
|
Instructions spécifiques:
|
|
1. Pour chaque image mentionnée: inclure la légende, source, page et description
|
|
2. Pour chaque tableau: inclure titre, source, page et signification
|
|
3. Pour les équations: utiliser la syntaxe LaTeX exacte
|
|
4. Ne pas inventer d'informations hors du contexte fourni
|
|
5. Citer précisément les sources
|
|
|
|
Historique de conversation:
|
|
{chat_history}
|
|
|
|
Contexte:
|
|
{context}
|
|
|
|
Question: {question}
|
|
|
|
Réponds de façon structurée en intégrant les images, tableaux et équations disponibles.
|
|
TA RÉPONSE DOIT ÊTRE UNIQUEMENT ET ENTIÈREMENT EN {language}. CETTE RÈGLE EST ABSOLUE.
|
|
""")
|
|
|
|
# Assurer que la langue est bien passée dans le format du prompt
|
|
selected_language = LANGUAGE_MAPPING.get(language, "français")
|
|
messages = prompt_template.format_messages(
|
|
chat_history=history_text,
|
|
context=context,
|
|
question=message,
|
|
language=selected_language
|
|
)
|
|
|
|
# 5. Créer un handler de streaming personnalisé
|
|
handler = GradioStreamingHandler()
|
|
|
|
# 6. Créer un modèle LLM avec notre handler
|
|
streaming_llm = ChatOllama(
|
|
model=rag_bot.llm.model,
|
|
base_url=rag_bot.llm.base_url,
|
|
streaming=True,
|
|
callbacks=[handler]
|
|
)
|
|
|
|
# 7. Lancer la génération dans un thread pour ne pas bloquer l'UI
|
|
def generate_response():
|
|
streaming_llm.invoke(messages)
|
|
|
|
thread = threading.Thread(target=generate_response)
|
|
thread.start()
|
|
|
|
# 8. Récupérer les tokens et mettre à jour l'interface
|
|
partial_response = ""
|
|
|
|
# Attendre les tokens avec un timeout
|
|
while thread.is_alive() or not handler.tokens_queue.empty():
|
|
try:
|
|
token = handler.tokens_queue.get(timeout=0.05)
|
|
partial_response += token
|
|
|
|
# Nettoyer la réponse uniquement pour l'affichage (pas pour l'historique interne)
|
|
clean_response = clean_llm_response(partial_response)
|
|
# Mettre à jour le dernier message (assistant)
|
|
messages_history[-1]["content"] = clean_response
|
|
yield messages_history, "", None, None
|
|
except queue.Empty:
|
|
continue
|
|
|
|
# Après la boucle, nettoyer la réponse complète pour l'historique interne
|
|
partial_response = clean_llm_response(partial_response)
|
|
rag_bot.chat_history.append({"role": "user", "content": message})
|
|
rag_bot.chat_history.append({"role": "assistant", "content": partial_response})
|
|
|
|
# 10. Récupérer les sources, images, tableaux
|
|
texts, images, tables = rag_bot._process_documents(docs)
|
|
|
|
# Préparer les informations sur les sources
|
|
source_info = ""
|
|
if texts:
|
|
source_info += f"📚 {len(texts)} textes • "
|
|
if images:
|
|
source_info += f"🖼️ {len(images)} images • "
|
|
if tables:
|
|
source_info += f"📊 {len(tables)} tableaux"
|
|
|
|
if source_info:
|
|
source_info = "Sources trouvées: " + source_info
|
|
|
|
# 11. Traiter les images
|
|
if show_sources and images:
|
|
images = images[:max_images]
|
|
for img in images:
|
|
img_data = img.get("image_data")
|
|
if img_data:
|
|
image = base64_to_image(img_data)
|
|
if image:
|
|
current_images.append({
|
|
"image": image,
|
|
"caption": img.get("caption", ""),
|
|
"source": img.get("source", ""),
|
|
"page": img.get("page", ""),
|
|
"description": img.get("description", "")
|
|
})
|
|
|
|
# 12. Traiter les tableaux
|
|
if show_sources and tables:
|
|
for table in tables:
|
|
current_tables.append({
|
|
"data": rag_bot.format_table(table.get("table_data", "")),
|
|
"caption": table.get("caption", ""),
|
|
"source": table.get("source", ""),
|
|
"page": table.get("page", ""),
|
|
"description": table.get("description", "")
|
|
})
|
|
|
|
# 13. Retourner les résultats finaux
|
|
images_display = display_images()
|
|
tables_display = display_tables()
|
|
yield messages_history, source_info, images_display, tables_display
|
|
|
|
else:
|
|
# Version sans streaming
|
|
print("Mode non-streaming activé")
|
|
source_info = ""
|
|
|
|
result = rag_bot.chat(
|
|
message,
|
|
stream=False,
|
|
language=LANGUAGE_MAPPING.get(language, "français") # Vérifiez que cette ligne existe
|
|
)
|
|
# Nettoyer la réponse des balises <think>
|
|
result["response"] = clean_llm_response(result["response"])
|
|
|
|
# Convertir l'historique au format messages
|
|
messages_history = convert_to_messages_format(history)
|
|
messages_history.append({"role": "user", "content": message})
|
|
messages_history.append({"role": "assistant", "content": result["response"]})
|
|
|
|
# Mise à jour de l'historique interne
|
|
rag_bot.chat_history.append({"role": "user", "content": message})
|
|
rag_bot.chat_history.append({"role": "assistant", "content": result["response"]})
|
|
|
|
# Traiter les sources
|
|
if "texts" in result:
|
|
source_info += f"📚 {len(result['texts'])} textes • "
|
|
if "images" in result:
|
|
source_info += f"🖼️ {len(result['images'])} images • "
|
|
if "tables" in result:
|
|
source_info += f"📊 {len(result['tables'])} tableaux"
|
|
|
|
if source_info:
|
|
source_info = "Sources trouvées: " + source_info
|
|
|
|
# Traiter les images et tableaux
|
|
if show_sources and "images" in result and result["images"]:
|
|
images = result["images"][:max_images]
|
|
for img in images:
|
|
img_data = img.get("image_data")
|
|
if img_data:
|
|
image = base64_to_image(img_data)
|
|
if image:
|
|
current_images.append({
|
|
"image": image,
|
|
"caption": img.get("caption", ""),
|
|
"source": img.get("source", ""),
|
|
"page": img.get("page", ""),
|
|
"description": img.get("description", "")
|
|
})
|
|
|
|
if show_sources and "tables" in result and result["tables"]:
|
|
tables = result["tables"]
|
|
for table in tables:
|
|
current_tables.append({
|
|
"data": rag_bot.format_table(table.get("table_data", "")),
|
|
"caption": table.get("caption", ""),
|
|
"source": table.get("source", ""),
|
|
"page": table.get("page", ""),
|
|
"description": table.get("description", "")
|
|
})
|
|
|
|
yield messages_history, source_info, display_images(), display_tables()
|
|
|
|
except Exception as e:
|
|
error_msg = f"Une erreur est survenue: {str(e)}"
|
|
traceback_text = traceback.format_exc()
|
|
print(error_msg)
|
|
print(traceback_text)
|
|
history = history + [(message, error_msg)]
|
|
yield history, "Erreur lors du traitement de la requête", None, None
|
|
|
|
# Fonction pour réinitialiser la conversation
|
|
def reset_conversation():
|
|
global current_images, current_tables
|
|
current_images = []
|
|
current_tables = []
|
|
|
|
rag_bot.clear_history()
|
|
|
|
# Retourner une liste vide au format messages
|
|
return [], "", None, None |