Upload 3 files

rag_chain/chain.py

@@ -13,57 +13,44 @@ from langchain_core.runnables import RunnableLambda
 from langchain_openai.chat_models import ChatOpenAI
 
 from .prompt_template import generate_prompt_template
-from .retrievers_setup import (DenseRetrieverClient, SparseRetrieverClient,
-                               compression_retriever_setup)
+from .retrievers_setup import (
+    DenseRetrieverClient,
+    SparseRetrieverClient,
+    compression_retriever_setup,
+    multi_query_retriever_setup,
+)
 
 # Helpers
 
 
 def reorder_documents(docs: list[Document]) -> list[Document]:
-    """Long-Context Reorder: No matter the architecture of the model, there is 
-    a performance degradation when we include 10+ retrieved documents.
-
-    Args:
-        docs (list): List of Langchain documents
-
-    Returns:
-        list: Reordered list of Langchain documents
-    """
-    reorder = LongContextReorder()
-    return reorder.transform_documents(docs)
+    """Reorder documents to mitigate performance degradation with long contexts."""
+    return LongContextReorder().transform_documents(docs)
 
 
 def randomize_documents(documents: list[Document]) -> list[Document]:
-    """Randomize the documents as an attempt to randomize the model recommendations."""
+    """Randomize documents to vary model recommendations."""
     random.shuffle(documents)
     return documents
 
 
-def format_practitioners_docs(docs: list[Document]) -> str:
-    """Format the practitioners_db Documents to markdown.
-    Args:
-        docs (list[Documents]): List of Langchain documents
-    Returns:
-        docs (str):
-    """
-    return f"\n{'-' * 3}\n".join(
-        [f"- Practitioner #{i+1}:\n\n\t" +
-            d.page_content for i, d in enumerate(docs)]
-    )
-
-
-def format_tall_tree_docs(docs: list[Document]) -> str:
-    """Format the tall_tree_db Documents to markdown.
-    Args:
-        docs (list[Documents]): List of Langchain documents
-    Returns:
-        docs (str):
+class DocumentFormatter:
+    def __init__(self, prefix: str):
+        self.prefix = prefix
 
-    """
-    return f"\n{'-' * 3}\n".join(
-        [f"- No. {i+1}:\n\n\t" +
-            d.page_content for i, d in enumerate(docs)]
-    )
+    def __call__(self, docs: list[Document]) -> str:
+        """Format the Documents to markdown.
+        Args:
+            docs (list[Documents]): List of Langchain documents
+        Returns:
+            docs (str):
+        """
+        return f"\n---\n".join(
+            [
+                f"- {self.prefix} {i+1}:\n\n\t" + d.page_content
+                for i, d in enumerate(docs)
+            ]
+        )
 
 
 @cache
@@ -74,8 +61,7 @@ def create_langsmith_client():
     os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"
     langsmith_api_key = os.getenv("LANGCHAIN_API_KEY")
     if not langsmith_api_key:
-        raise EnvironmentError(
-            "Missing environment variable: LANGCHAIN_API_KEY")
+        raise EnvironmentError("Missing environment variable: LANGCHAIN_API_KEY")
     return langsmith.Client()
 
 
@@ -83,7 +69,9 @@ def create_langsmith_client():
 
 
 @cache
-def get_rag_chain(model_name: str = "gpt-4", temperature: float = 0.2) -> tuple[ChatOpenAI, ConversationBufferWindowMemory]:
+def get_rag_chain(
+    model_name: str = "gpt-4", temperature: float = 0.2
+) -> tuple[ChatOpenAI, ConversationBufferWindowMemory]:
     """Set up runnable and chat memory
 
     Args:
@@ -94,78 +82,105 @@ def get_rag_chain(model_name: str = "gpt-4", temperature: float = 0.2) -> tuple[
         Runnable, Memory: Chain and Memory
     """
 
+    RETRIEVER_PARAMETERS = {
+        "embeddings_model": "text-embedding-3-small",
+        "k_dense_practitioners_db": 8,
+        "k_sparse_practitioners_db": 15,
+        "weights_ensemble_practitioners_db": [0.2, 0.8],
+        "k_compression_practitioners_db": 18,
+        "k_dense_talltree": 6,
+        "k_compression_talltree": 6,
+    }
+
     # Set up Langsmith to trace the chain
     langsmith_tracing = create_langsmith_client()
 
     # LLM and prompt template
-    llm = ChatOpenAI(model_name=model_name,
-                     temperature=temperature)
+    llm = ChatOpenAI(
+        model_name=model_name,
+        temperature=temperature,
+    )
 
     prompt = generate_prompt_template()
 
     # Set retrievers pointing to the practitioners's dataset
-    embeddings_model = "text-embedding-ada-002"
-    dense_retriever_client = DenseRetrieverClient(embeddings_model=embeddings_model,
-                                                  collection_name="practitioners_db")
+    dense_retriever_client = DenseRetrieverClient(
+        embeddings_model=RETRIEVER_PARAMETERS["embeddings_model"],
+        collection_name="practitioners_db",
+        search_type="similarity",
+        k=RETRIEVER_PARAMETERS["k_dense_practitioners_db"],
+    )  # k x 4 using multiquery retriever
 
     # Qdrant db as a retriever
-    practitioners_db_dense_retriever = dense_retriever_client.get_dense_retriever(search_type="similarity",
-                                                                                  k=10)
+    practitioners_db_dense_retriever = dense_retriever_client.get_dense_retriever()
 
-    # Testing the sparse vector retriever using Qdrant
-    collection_name = "practitioners_db_sparse_collection"
-    vector_name = "sparse_vector"
+    # Multiquery retriever using the dense retriever
+    practitioners_db_dense_multiquery_retriever = multi_query_retriever_setup(
+        practitioners_db_dense_retriever
+    )
+
+    # Sparse vector retriever
     sparse_retriever_client = SparseRetrieverClient(
-        collection_name=collection_name,
-        vector_name=vector_name,
+        collection_name="practitioners_db_sparse_collection",
+        vector_name="sparse_vector",
         splade_model_id="naver/splade-cocondenser-ensembledistil",
-        k=15)
+        k=RETRIEVER_PARAMETERS["k_sparse_practitioners_db"],
+    )
+
     practitioners_db_sparse_retriever = sparse_retriever_client.get_sparse_retriever()
 
     # Ensemble retriever for hyprid search (dense retriever seems to work better but the dense retriever is good for acronyms like RMT)
     practitioners_ensemble_retriever = EnsembleRetriever(
-        retrievers=[practitioners_db_dense_retriever,
-                    practitioners_db_sparse_retriever], weights=[0.2, 0.8]
+        retrievers=[
+            practitioners_db_dense_multiquery_retriever,
+            practitioners_db_sparse_retriever,
+        ],
+        weights=RETRIEVER_PARAMETERS["weights_ensemble_practitioners_db"],
     )
 
     # Compression retriever for practitioners db
     practitioners_db_compression_retriever = compression_retriever_setup(
         practitioners_ensemble_retriever,
-        embeddings_model="text-embedding-ada-002",
-        similarity_threshold=0.74
+        embeddings_model=RETRIEVER_PARAMETERS["embeddings_model"],
+        k=RETRIEVER_PARAMETERS["k_compression_practitioners_db"],
     )
 
     # Set retrievers pointing to the tall_tree_db
-    dense_retriever_client = DenseRetrieverClient(embeddings_model=embeddings_model,
-                                                  collection_name="tall_tree_db")
-    tall_tree_db_dense_retriever = dense_retriever_client.get_dense_retriever(search_type="similarity",
-                                                                              k=8)
+    dense_retriever_client = DenseRetrieverClient(
+        embeddings_model=RETRIEVER_PARAMETERS["embeddings_model"],
+        collection_name="tall_tree_db",
+        search_type="similarity",
+        k=RETRIEVER_PARAMETERS["k_dense_talltree"],
+    )
+
+    tall_tree_db_dense_retriever = dense_retriever_client.get_dense_retriever()
+
     # Compression retriever for tall_tree_db
     tall_tree_db_compression_retriever = compression_retriever_setup(
         tall_tree_db_dense_retriever,
-        embeddings_model="text-embedding-ada-002",
-        similarity_threshold=0.74
+        embeddings_model=RETRIEVER_PARAMETERS["embeddings_model"],
+        k=RETRIEVER_PARAMETERS["k_compression_talltree"],
     )
 
     # Set conversation history window memory. It only uses the last k interactions.
-    memory = ConversationBufferWindowMemory(memory_key="history",
-                                            return_messages=True,
-                                            k=6)
+    memory = ConversationBufferWindowMemory(
+        memory_key="history",
+        return_messages=True,
+        k=6,
+    )
 
     # Set up runnable using LCEL
-    setup_and_retrieval = {"practitioners_db": itemgetter("message")
-                           | practitioners_db_compression_retriever
-                           | format_practitioners_docs,
-                           "tall_tree_db": itemgetter("message") | tall_tree_db_compression_retriever | format_tall_tree_docs,
-                           "history": RunnableLambda(memory.load_memory_variables) | itemgetter("history"),
-                           "message": itemgetter("message")
-                           }
-
-    chain = (
-        setup_and_retrieval
-        | prompt
-        | llm
-        | StrOutputParser()
-    )
+    setup_and_retrieval = {
+        "practitioners_db": itemgetter("message")
+        | practitioners_db_compression_retriever
+        | DocumentFormatter("Practitioner #"),
+        "tall_tree_db": itemgetter("message")
+        | tall_tree_db_compression_retriever
+        | DocumentFormatter("No."),
+        "history": RunnableLambda(memory.load_memory_variables) | itemgetter("history"),
+        "message": itemgetter("message"),
+    }
+
+    chain = setup_and_retrieval | prompt | llm | StrOutputParser()
 
     return chain, memory

rag_chain/prompt_template.py

@@ -1,7 +1,8 @@
-from langchain.prompts import (ChatPromptTemplate,
-                               SystemMessagePromptTemplate,
-                               MessagesPlaceholder,
-                               )
+from langchain.prompts import (
+    ChatPromptTemplate,
+    SystemMessagePromptTemplate,
+    MessagesPlaceholder,
+)
 
 
 def generate_prompt_template():
@@ -83,8 +84,7 @@ You are a helpful Virtual Assistant at Tall Tree Health in British Columbia, Can
 """
 
     # Template for system message with markdown formatting
-    system_message = SystemMessagePromptTemplate.from_template(
-        system_template)
+    system_message = SystemMessagePromptTemplate.from_template(system_template)
 
     prompt = ChatPromptTemplate.from_messages(
         [

rag_chain/retrievers_setup.py

@@ -14,139 +14,144 @@ from langchain_openai.embeddings import OpenAIEmbeddings
 from transformers import AutoModelForMaskedLM, AutoTokenizer
 
 
-class DenseRetrieverClient:
-    """Inititalize the dense retriever using OpenAI text embeddings and Qdrant vector database.
+class ValidateQdrantClient:
+    """Base class for retriever clients to ensure environment variables are set."""
 
-    Attributes:
-        embeddings_model (str): The embeddings model to use. Right now only OpenAI text embeddings.
-        collection_name (str): Qdrant collection name.
-        client (QdrantClient): Qdrant client.
-        qdrant_collection (Qdrant): Qdrant collection.
-    """
-
-    def __init__(self, embeddings_model: str = "text-embedding-ada-002", collection_name: str = "practitioners_db"):
+    def __init__(self):
         self.validate_environment_variables()
+
+    def validate_environment_variables(self):
+        """Check if the Qdrant environment variables are set."""
+        required_vars = ["QDRANT_API_KEY", "QDRANT_URL"]
+        missing_vars = [var for var in required_vars if not os.getenv(var)]
+        if missing_vars:
+            raise EnvironmentError(
+                f"Missing environment variable(s): {', '.join(missing_vars)}"
+            )
+
+
+class DenseRetrieverClient(ValidateQdrantClient):
+    """Initialize the dense retriever using OpenAI text embeddings and Qdrant vector database."""
+
+    TEXT_EMBEDDING_MODELS = [
+        "text-embedding-ada-002",
+        "text-embedding-3-small",
+        "text-embedding-3-large",
+    ]
+
+    def __init__(
+        self,
+        embeddings_model="text-embedding-3-small",
+        collection_name="practitioners_db",
+        search_type="similarity",
+        k=4,
+    ):
+        super().__init__()
+        if embeddings_model not in self.TEXT_EMBEDDING_MODELS:
+            raise ValueError(
+                f"Invalid embeddings model: {embeddings_model}. Valid options are {', '.join(self.TEXT_EMBEDDING_MODELS)}."
+            )
         self.embeddings_model = embeddings_model
         self.collection_name = collection_name
+        self.search_type = search_type
+        self.k = k
         self.client = qdrant_client.QdrantClient(
             url=os.getenv("QDRANT_URL"),
             api_key=os.getenv("QDRANT_API_KEY"),
             prefer_grpc=True,
         )
-        self.qdrant_collection = self.load_qdrant_collection()
-
-    def validate_environment_variables(self):
-        """ Check if the Qdrant environment variables are set."""
-        required_vars = ["QDRANT_API_KEY", "QDRANT_URL"]
-        for var in required_vars:
-            if not os.getenv(var):
-                raise EnvironmentError(f"Missing environment variable: {var}")
+        self._qdrant_collection = None
 
     def set_qdrant_collection(self, embeddings):
         """Prepare the Qdrant collection for the embeddings model."""
-        return Qdrant(client=self.client,
-                      collection_name=self.collection_name,
-                      embeddings=embeddings)
+        return Qdrant(
+            client=self.client,
+            collection_name=self.collection_name,
+            embeddings=embeddings,
+        )
 
+    @property
     @cache
-    def load_qdrant_collection(self):
+    def qdrant_collection(self):
         """Load Qdrant collection for a given embeddings model."""
-        # TODO: Test new OpenAI text embeddings models
-        openai_text_embeddings = ["text-embedding-ada-002"]
-        if self.embeddings_model in openai_text_embeddings:
-            self.qdrant_collection = self.set_qdrant_collection(
-                OpenAIEmbeddings(model=self.embeddings_model))
-        else:
-            raise ValueError(
-                f"Invalid embeddings model: {self.embeddings_model}. Valid options are {', '.join(openai_text_embeddings)}.")
-
-        return self.qdrant_collection
-
-    def get_dense_retriever(self, search_type: str = "similarity", k: int = 4):
-        """Set up retrievers (Qdrant vectorstore as retriever).
-
-        Args:
-            search_type (str, optional): similarity or mmr. Defaults to "similarity".
-            k (int, optional): Number of documents retrieved. Defaults to 4.
-
-        Returns:
-            Retriever: Vectorstore as a retriever
-        """
-        dense_retriever = self.qdrant_collection.as_retriever(search_type=search_type,
-                                                              search_kwargs={
-                                                                  "k": k}
-                                                              )
-        return dense_retriever
+        if self._qdrant_collection is None:
+            self._qdrant_collection = self.set_qdrant_collection(
+                OpenAIEmbeddings(model=self.embeddings_model)
+            )
+        return self._qdrant_collection
+
+    def get_dense_retriever(self):
+        """Set up retrievers (Qdrant vectorstore as retriever)."""
+        return self.qdrant_collection.as_retriever(
+            search_type=self.search_type, search_kwargs={"k": self.k}
+        )
 
 
-class SparseRetrieverClient:
-    """Inititalize the sparse retriever using the SPLADE neural retrieval model and Qdrant vector database.
+class SparseRetrieverClient(ValidateQdrantClient):
+    """Initialize the sparse retriever using the SPLADE neural retrieval model and Qdrant vector database."""
 
-    Attributes:
-        collection_name (str): Qdrant collection name.
-        vector_name (str): Qdrant vector name.
-        splade_model_id (str): The SPLADE neural retrieval model id.
-        k (int): Number of documents retrieved.
-        client (QdrantClient): Qdrant client.
-    """
+    def __init__(
+        self,
+        collection_name,
+        vector_name,
+        splade_model_id="naver/splade-cocondenser-ensembledistil",
+        k=15,
+    ):
 
-    def __init__(self, collection_name: str, vector_name: str, splade_model_id: str = "naver/splade-cocondenser-ensembledistil", k: int = 15):
-        self.validate_environment_variables()
-        self.client = qdrant_client.QdrantClient(url=os.getenv(
-            "QDRANT_URL"), api_key=os.getenv("QDRANT_API_KEY"))
+        # Validate Qdrant client
+        super().__init__()
+        self.client = qdrant_client.QdrantClient(
+            url=os.getenv("QDRANT_URL"), api_key=os.getenv("QDRANT_API_KEY")
+        )  # TODO: prefer_grpc=True is not working
         self.model_id = splade_model_id
-        self.tokenizer, self.model = self.set_tokenizer_config()
+        self._tokenizer = None
+        self._model = None
         self.collection_name = collection_name
         self.vector_name = vector_name
         self.k = k
 
-    def validate_environment_variables(self):
-        required_vars = ["QDRANT_API_KEY", "QDRANT_URL"]
-        for var in required_vars:
-            if not os.getenv(var):
-                raise EnvironmentError(f"Missing environment variable: {var}")
+    @property
+    @cache
+    def tokenizer(self):
+        """Initialize the tokenizer."""
+        if self._tokenizer is None:
+            self._tokenizer = AutoTokenizer.from_pretrained(self.model_id)
+        return self._tokenizer
 
+    @property
     @cache
-    def set_tokenizer_config(self):
-        """Initialize the tokenizer and the SPLADE neural retrieval model.
-        See to https://huggingface.co/naver/splade-cocondenser-ensembledistil for more details.
-        """
-        tokenizer = AutoTokenizer.from_pretrained(self.model_id)
-        model = AutoModelForMaskedLM.from_pretrained(self.model_id)
-        return tokenizer, model
+    def model(self):
+        """Initialize the SPLADE neural retrieval model."""
+        if self._model is None:
+            self._model = AutoModelForMaskedLM.from_pretrained(self.model_id)
+        return self._model
 
     def sparse_encoder(self, text: str) -> tuple[list[int], list[float]]:
-        """This function encodes the input text into a sparse vector. The encoder is required for the QdrantSparseVectorRetriever.
-        Adapted from the Qdrant documentation: Computing the Sparse Vector code.
-
-        Args:
-            text (str): Text to encode
-
-        Returns:
-            tuple[list[int], list[float]]: Indices and values of the sparse vector
-        """
-        tokens = self.tokenizer(text, return_tensors="pt",
-                                max_length=512, padding="max_length", truncation=True)
+        """Encode the input text into a sparse vector."""
+        tokens = self.tokenizer(
+            text,
+            return_tensors="pt",
+            max_length=512,
+            padding="max_length",
+            truncation=True,
+        )
 
         with torch.no_grad():
-            output = self.model(**tokens)
-
-        logits, attention_mask = output.logits, tokens.attention_mask
+            logits = self.model(**tokens).logits
 
         relu_log = torch.log1p(torch.relu(logits))
-        weighted_log = relu_log * attention_mask.unsqueeze(-1)
-
-        max_val, _ = torch.max(weighted_log, dim=1)
-        vec = max_val.squeeze()
+        weighted_log = relu_log * tokens.attention_mask.unsqueeze(-1)
 
-        indices = torch.nonzero(vec, as_tuple=False).squeeze().cpu().numpy()
-        values = vec[indices].cpu().numpy()
+        max_val = torch.max(weighted_log, dim=1).values.squeeze()
+        indices = torch.nonzero(max_val, as_tuple=False).squeeze().cpu().numpy()
+        values = max_val[indices].cpu().numpy()
 
         return indices.tolist(), values.tolist()
 
-    def get_sparse_retriever(self):
+    def get_sparse_retriever(self) -> QdrantSparseVectorRetriever:
+        """Return a Qdrant vector sparse retriever."""
 
-        sparse_retriever = QdrantSparseVectorRetriever(
+        return QdrantSparseVectorRetriever(
             client=self.client,
             collection_name=self.collection_name,
             sparse_vector_name=self.vector_name,
@@ -154,63 +159,38 @@ class SparseRetrieverClient:
             k=self.k,
         )
 
-        return sparse_retriever
-
-
-def compression_retriever_setup(base_retriever, embeddings_model: str = "text-embedding-ada-002", similarity_threshold: float = 0.76) -> ContextualCompressionRetriever:
-    """
-    Creates a ContextualCompressionRetriever with a base retriever and a similarity threshold.
-
-    The ContextualCompressionRetriever uses an EmbeddingsFilter with OpenAIEmbeddings to filter out documents 
-    with a similarity score below the given threshold.
 
-    Args:
-        base_retriever: Retriever to be filtered.
-        similarity_threshold (float, optional): The similarity threshold for the EmbeddingsFilter. 
-            Documents with a similarity score below this threshold will be filtered out. Defaults to 0.76 (Obtained by experimenting with text-embeddings-ada-002).
-            ** Be careful with this parameter, as it can have a big impact on the results and highly depends on the embeddings model used.
+def compression_retriever_setup(
+    base_retriever, embeddings_model="text-embedding-3-small", k=20
+):
+    """Creates a ContextualCompressionRetriever with an EmbeddingsFilter."""
+    filter = EmbeddingsFilter(embeddings=OpenAIEmbeddings(model=embeddings_model), k=k)
 
-    Returns:
-        ContextualCompressionRetriever: The created ContextualCompressionRetriever.
-    """
-
-    # Set up compression retriever (filter out documents with low similarity)
-    relevant_filter = EmbeddingsFilter(embeddings=OpenAIEmbeddings(model=embeddings_model),
-                                       similarity_threshold=similarity_threshold)
-
-    compression_retriever = ContextualCompressionRetriever(
-        base_compressor=relevant_filter, base_retriever=base_retriever
+    return ContextualCompressionRetriever(
+        base_compressor=filter, base_retriever=base_retriever
     )
 
-    return compression_retriever
-
-
-def multi_query_retriever_setup(retriever) -> MultiQueryRetriever:
-    """ Configure a multi-query retriever using a base retriever and the LLM.
-
-    Args:
-        retriever:
 
-    Returns:
-        retriever: MultiQueryRetriever
-    """
+def multi_query_retriever_setup(retriever):
+    """Configure a multi-query retriever using a base retriever."""
 
-    QUERY_PROMPT = PromptTemplate(
+    prompt = PromptTemplate(
         input_variables=["question"],
         template="""
-
-        Your task is to generate 3 different versions of the provided question, incorporating the user's location preference in each version. Each version must be separated by newlines. Ensure that no part of your response is enclosed in quotation marks. Do not modify any acronyms or unfamiliar terms. Keep your responses clear, concise, and limited to these alternatives. 
-        Note: The text provided are queries to Tall Tree Health Centre's AI virtual assistant.
-
-        Question:
+        
+        Your task is to generate 3 different grammatically correct versions of the provided text,
+        incorporating the user's location preference in each version. Format these versions as paragraphs and present them as items in a Markdown formatted numbered list ("1. "). There should be no additional new lines or spaces between each version. Do not enclose your response in quotation marks. Do not modify unfamiliar acronyms and keep your responses clear and concise.
+        
+        **Notes**: The text provided are user questions to Tall Tree Health Centre's AI virtual assistant. `Location preference:` is the location of the Tall Tree Health clinic that the user prefers.
+
+        Text to be modified:
+        ```
         {question}
-
-        """,
+        ```""",
     )
 
-    llm = ChatOpenAI(model='gpt-3.5-turbo', temperature=0)
-    multi_query_retriever = MultiQueryRetriever.from_llm(
-        retriever=retriever, llm=llm, prompt=QUERY_PROMPT, include_original=True,
-    )
+    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
 
-    return multi_query_retriever
+    return MultiQueryRetriever.from_llm(
+        retriever=retriever, llm=llm, prompt=prompt, include_original=True
+    )