Refactor RAG system query methods; update descriptions and improve logging for better clarity

app.py

@@ -12,7 +12,7 @@ class ChatInterface:
 
     def respond(self, message: str, history: List[List[str]]):
         result = ""
-        for text in self.rag_system.query_iter(message, history):
+        for text in self.rag_system.query(message, history):
             result += text
             yield result
         return result
@@ -21,7 +21,7 @@ class ChatInterface:
         chat_interface = gr.ChatInterface(
             fn=self.respond,
             title="Medivocate",
-            description="Medivocate is an AI-driven platform leveraging Retrieval-Augmented Generation (RAG) powered by African history. It processes and classifies document pages with precision to provide trustworthy, personalized guidance, fostering accurate knowledge and equitable access to historical insights.",
+            description="Medivocate est une application qui offre des informations claires et structurées sur l'histoire de l'Afrique et sa médecine traditionnelle, en s'appuyant exclusivement sur un contexte issu de documentaires sur l'histoire du continent africain.",
             # retry_btn=None,
             # undo_btn=None,
             # clear_btn="Clear",

src/rag_pipeline/prompts.py

@@ -0,0 +1,43 @@
+from langchain.prompts.chat import (
+    ChatPromptTemplate,
+    HumanMessagePromptTemplate,
+    MessagesPlaceholder,
+    SystemMessagePromptTemplate,
+)
+
+system_template = """
+Vous êtes un assistant IA qui fournit des informations sur l'histoire de l'Afrique et la médecine traditionnelle africaine. Vous recevez une question et fournissez une réponse claire et structurée. Lorsque cela est pertinent, utilisez des points et des listes pour structurer vos réponses.
+
+Utilisez uniquement les éléments de contexte suivants pour répondre à la question de l'utilisateur. Si vous ne connaissez pas la réponse, dites simplement que vous ne savez pas, n'essayez pas d'inventer une réponse.
+
+Si la question posée est dans une langue parlée en Afrique ou demande une traduction dans une de ces langues, répondez que vous ne savez pas et demandez à l'utilisateur de reformuler sa question.
+
+Si vous connaissez la réponse à la question mais que cette réponse ne provient pas du contexte ou n'est pas relative à l'histoire africaine ou à la médecine traditionnelle, répondez que vous ne savez pas et demandez à l'utilisateur de reformuler sa question.
+
+-----------------
+{context}
+"""
+
+messages = [
+    MessagesPlaceholder(variable_name="chat_history"),
+    SystemMessagePromptTemplate.from_template(system_template),
+    HumanMessagePromptTemplate.from_template("{input}"),
+]
+CHAT_PROMPT = ChatPromptTemplate.from_messages(messages)
+
+
+contextualize_q_system_prompt = (
+    "Étant donné un historique de conversation et la dernière question de l'utilisateur "
+    "qui pourrait faire référence au contexte dans l'historique de conversation, "
+    "formulez une question autonome qui peut être comprise "
+    "sans l'historique de conversation. NE répondez PAS à la question, reformulez-la "
+    "si nécessaire, sinon retournez-la telle quelle."
+)
+
+CONTEXTUEL_QUERY_PROMPT = ChatPromptTemplate.from_messages(
+    [
+        SystemMessagePromptTemplate.from_template(contextualize_q_system_prompt),
+        MessagesPlaceholder("chat_history"),
+        HumanMessagePromptTemplate.from_template("{input}"),
+    ]
+)

src/rag_pipeline/rag_system.py

@@ -1,13 +1,18 @@
-import os
+import logging
 from typing import List, Optional
 
 from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain.chains.conversational_retrieval.base import (
+    ConversationalRetrievalChain,
+)
+from langchain.chains.history_aware_retriever import (
+    create_history_aware_retriever,
+)
 from langchain.chains.retrieval import create_retrieval_chain
-from langchain.prompts import PromptTemplate
-from langchain_core.runnables import Runnable
 
-from ..utilities.llm_models import get_llm_model_chat, get_llm_model_embedding
+from ..utilities.llm_models import get_llm_model_chat
 from ..vector_store.vector_store import VectorStoreManager
+from .prompts import CHAT_PROMPT, CONTEXTUEL_QUERY_PROMPT
 
 
 class RAGSystem:
@@ -19,19 +24,16 @@ class RAGSystem:
         top_k_documents=5,
     ):
         self.top_k_documents = top_k_documents
-        self.embeddings = self._get_embeddings()
         self.llm = self._get_llm()
-        self.chain: Optional[Runnable] = None
+        self.chain: Optional[ConversationalRetrievalChain] = None
         self.vector_store_management = VectorStoreManager(
             docs_dir, persist_directory_dir, batch_size
         )
 
-    def _get_llm(self):
-        return get_llm_model_chat("GROQ", temperature=0.1, max_tokens=500)
-
-    def _get_embeddings(self):
-        """Initialize embeddings based on environment configuration"""
-        return get_llm_model_embedding()
+    def _get_llm(
+        self,
+    ):
+        return get_llm_model_chat(temperature=0.1, max_tokens=1000)
 
     def initialize_vector_store(self, documents: List = None):
         """Initialize or load the vector store"""
@@ -40,79 +42,28 @@ class RAGSystem:
     def setup_rag_chain(self):
         if self.chain is not None:
             return
-        """Set up the RAG chain with custom prompt"""
-        prompt_template = """Inspirez vous du contexte fourni ci-dessous pour répondre à la question qui suit de la manière la plus précise possible.  
-Si la réponse ne peut pas être déterminée à partir du contexte, évitez d'inventer des informations.
-L'historique ici fait référence aux précédents échanges avec un utilisateur, tu devrais l'ignore si aucun rapport avec la question posée.
-Tes réponses doivent être naturelles sous forme de faits, au lieu de faire mention du fait que réponds en fonction d'un contexte.
-
-**Historique** :
-{history}
-
-**Contexte** : 
-{context}
-
-**Question** :
-{input}
-
-Réponse (Vous devez répondre dans la même langue que celle de la question) :"""
-
-        prompt = PromptTemplate(
-            template=prompt_template, input_variables=["context", "input", "history"]
-        )
         retriever = self.vector_store_management.vector_store.as_retriever(
             search_kwargs={"k": self.top_k_documents}
         )
-        question_answer_chain = create_stuff_documents_chain(self.llm, prompt)
-
-        self.chain = create_retrieval_chain(retriever, question_answer_chain)
-
-    def query(self, question: str, history: List[tuple[str]] = []):
-        """Query the RAG system"""
-        if not self.vector_store_management.vector_store:
-            self.initialize_vector_store()
 
-        self.setup_rag_chain()
-
-        # Format history as a single string of interactions
-        history_text = "\n".join(
-            [f"Utilisateur: {user}\nAssistant: {assistant}" for user, assistant in history]
+        # Contextualize question
+        history_aware_retriever = create_history_aware_retriever(
+            self.llm, retriever, CONTEXTUEL_QUERY_PROMPT
         )
+        question_answer_chain = create_stuff_documents_chain(self.llm, CHAT_PROMPT)
+        self.chain = create_retrieval_chain(
+            history_aware_retriever, question_answer_chain
+        )
+        logging.info("RAG chain setup complete" + str(self.chain))
+        return self.chain
 
-        response = self.chain.invoke({"input": question,  "history": history_text})
-
-        return {
-            "answer": response["answer"],
-            "source_documents": [doc.page_content for doc in response["context"]],
-        }
-
-    def query_iter(self, question: str, history: List[tuple[str]] = []):
+    def query(self, question: str, history: list = []):
         """Query the RAG system"""
         if not self.vector_store_management.vector_store:
             self.initialize_vector_store()
 
         self.setup_rag_chain()
 
-        # Format history as a single string of interactions
-        history_text = "\n".join(
-            [f"Utilisateur: {user}\nAssistant: {assistant}" for user, assistant in history]
-        )
-
-        for token in self.chain.stream({"input": question, "history": history_text}):
+        for token in self.chain.stream({"input": question, "chat_history": history}):
             if "answer" in token:
                 yield token["answer"]
-
-
-if __name__ == "__main__":
-    from glob import glob
-
-    docs_dir = "data/docs"
-    persist_directory_dir = "data/chroma_db"
-    batch_size = 64
-
-    # Initialize RAG system
-    rag = RAGSystem(docs_dir, persist_directory_dir, batch_size)
-
-    rag.initialize_vector_store()  # vector store initialized
-
-    print(rag.query("Quand a eu lieu la traite négrière ?"))

src/utilities/llm_models.py

@@ -1,10 +1,9 @@
 import os
 from enum import Enum
-from typing import Union
 
 from langchain_groq import ChatGroq
-from langchain_ollama import ChatOllama, OllamaEmbeddings
 from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_ollama import ChatOllama, OllamaEmbeddings
 
 
 class LLMModel(Enum):
@@ -12,13 +11,11 @@ class LLMModel(Enum):
     GROQ = ChatGroq
 
 
-def get_llm_model_chat(
-    model_type: Union[str, LLMModel], temperature=0, max_tokens=None
-):
-    if isinstance(model_type, str):
-        model_type = LLMModel[model_type.upper()]
-    if model_type == LLMModel.OLLAMA:
-        return model_type.value(
+def get_llm_model_chat(temperature=0.01, max_tokens=None):
+    if str(os.getenv("USE_OLLAMA_CHAT")) == "1" and "localhost" not in str(
+        os.getenv("OLLAMA_HOST")
+    ):
+        return ChatOllama(
             model=os.getenv("OLLAMA_MODEL"),
             temperature=temperature,
             max_tokens=max_tokens,
@@ -30,7 +27,7 @@ def get_llm_model_chat(
                 }
             },
         )
-    return model_type.value(
+    return ChatGroq(
         model=os.getenv("GROQ_MODEL_NAME"),
         temperature=temperature,
         max_tokens=max_tokens,
@@ -38,7 +35,7 @@ def get_llm_model_chat(
 
 
 def get_llm_model_embedding():
-    if os.getenv("USE_HF"):
+    if str(os.getenv("USE_HF_EMBEDDING")) == "1":
         return HuggingFaceEmbeddings(
             model_name=os.getenv("HF_MODEL"),  # You can replace with any HF model
             model_kwargs={"device": "cpu"},

src/vector_store/vector_store.py

@@ -1,33 +1,11 @@
-import json
 import os
-from concurrent.futures import ThreadPoolExecutor
-from glob import glob
 from typing import List
 
-from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_chroma import Chroma
-from langchain_community.document_loaders import DirectoryLoader, TextLoader
-from langchain_core.documents import Document
 from tqdm import tqdm
 
 from ..utilities.llm_models import get_llm_model_embedding
 
-
-def sanitize_metadata(metadata):
-    sanitized = {}
-    for key, value in metadata.items():
-        if isinstance(value, list):
-            # Convert lists to comma-separated strings or handle appropriately
-            sanitized[key] = ", ".join(value)
-        elif isinstance(value, (str, int, float, bool)):
-            sanitized[key] = value
-        else:
-            raise ValueError(
-                f"Unsupported metadata type for key '{key}': {type(value)}"
-            )
-    return sanitized
-
-
 class VectorStoreManager:
     def __init__(self, docs_dir: str, persist_directory_dir: str, batch_size=64):
         self.embeddings = get_llm_model_embedding()