Upload 4 files

rag_chain/__init__.py

@@ -0,0 +1 @@
+

rag_chain/chain.py

@@ -0,0 +1,173 @@
+import os
+import random
+from functools import cache
+from operator import itemgetter
+
+import langsmith
+from langchain.memory import ConversationBufferWindowMemory
+from langchain.retrievers import EnsembleRetriever
+from langchain_community.document_transformers import LongContextReorder
+from langchain_core.documents import Document
+from langchain_core.output_parsers import StrOutputParser
+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)
+
+# 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)
+
+
+def randomize_documents(documents: list[Document]) -> list[Document]:
+    """Randomize the documents to vary the 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):
+
+    """
+    return f"\n{'-' * 3}\n".join(
+        [f"- No. {i+1}:\n\n\t" +
+            d.page_content for i, d in enumerate(docs)]
+    )
+
+
+def create_langsmith_client():
+    """Create a Langsmith client."""
+    os.environ["LANGCHAIN_TRACING_V2"] = "true"
+    os.environ["LANGCHAIN_PROJECT"] = "talltree-ai-assistant"
+    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")
+    return langsmith.Client()
+
+
+# Set up Runnable and Memory
+
+
+@cache
+def get_rag_chain(model_name: str = "gpt-4", temperature: float = 0.2) -> tuple[ChatOpenAI, ConversationBufferWindowMemory]:
+    """Set up runnable and chat memory
+
+    Args:
+        model_name (str, optional): LLM model. Defaults to "gpt-4" 30012024.
+        temperature (float, optional): Model temperature. Defaults to 0.2.
+
+    Returns:
+        Runnable, Memory: Chain and Memory
+    """
+
+    # Set up Langsmith to trace the chain
+    langsmith_tracing = create_langsmith_client()
+
+    # LLM and prompt template
+    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")
+
+    # Qdrant db as a retriever
+    practitioners_db_dense_retriever = dense_retriever_client.get_dense_retriever(search_type="similarity",
+                                                                                  k=10)
+
+    # Testing the sparse vector retriever using Qdrant
+    collection_name = "practitioners_db_sparse_collection"
+    vector_name = "sparse_vector"
+    sparse_retriever_client = SparseRetrieverClient(
+        collection_name=collection_name,
+        vector_name=vector_name,
+        splade_model_id="naver/splade-cocondenser-ensembledistil",
+        k=15)
+    practitioners_db_sparse_retriever = sparse_retriever_client.get_sparse_retriever()
+
+    # TODO Test the ensemble retriever for hyprid search (dense retriever seems to work better)
+    # Using only the filtered sparse retriever
+    practitioners_ensemble_retriever = EnsembleRetriever(
+        retrievers=[practitioners_db_dense_retriever,
+                    practitioners_db_sparse_retriever], weights=[0.1, 0.9]
+    )
+
+    # Compression retriever for practitioners db
+    # TODO Test the filtered ensemble retriever *** Using only the sparse retriever ***
+    practitioners_db_compression_retriever = compression_retriever_setup(
+        practitioners_db_sparse_retriever,
+        embeddings_model="text-embedding-ada-002",
+        similarity_threshold=0.74
+    )
+
+    # 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=5)
+    # 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.5
+    )
+
+    # Set conversation history window memory. It only uses the last k=4 interactions.
+    memory = ConversationBufferWindowMemory(memory_key="history",
+                                            return_messages=True,
+                                            k=5)
+
+    # Set up runnable using LCEL
+    setup_and_retrieval = {"practitioners_db": itemgetter("message")
+                           | practitioners_db_compression_retriever
+                           | randomize_documents
+                           | 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()
+    )
+
+    return chain, memory

rag_chain/prompt_template.py

@@ -0,0 +1,102 @@
+from langchain.prompts import (ChatPromptTemplate,
+                               SystemMessagePromptTemplate,
+                               MessagesPlaceholder,
+                               )
+
+
+def generate_prompt_template():
+
+    # Prompt templates
+    system_template = """# Role
+
+---
+
+You are a helpful Virtual Assistant at Tall Tree Health in British Columbia, Canada. Your role is to analyze the patient's symptoms or needs and connect them with the appropriate practitioner or service offered by Tall Tree. Respond to Patient Queries using the `Practitioners Database` and `Tall Tree Health Centre Information` provided in the `Context`. Follow the `Response Guidelines` listed below:
+
+---
+
+# Response Guidelines
+
+1. **Interaction**: Engage in a warm, empathetic, and professional manner. Keep responses brief and focused on the patient's query. Use markdown formatting.
+
+2. **Symptoms/needs and Location Preference**: Ask for symptoms/needs and location preference (Cordova Bay, James Bay, and Vancouver) before recommending a practitioner or service.
+
+3. **Avoid Making Assumptions**: Stick to the given `Context`. If you're unable to assist, offer the user the contact details for the closest `Tall Tree Health` clinic.
+
+4. **No Medical Advice**: Refrain from giving any medical advice or acting as a healthcare professional. Avoid discussing healthcare costs.
+
+5. **Symptoms/needs and Service Verification**: Match the patient's symptoms/needs with the `Focus Area` field in the `Practitioners Database`. If no match is found, advise the patient accordingly without recommending a practitioner, as Tall Tree is not a primary healthcare provider.
+
+6. **Recommending Practitioners**: Use the patient's symptoms/needs and location to recommend 3 practitioners from the `Practitioners Database`. Focus on `Discipline`, `Focus Areas`, `Location`, `Treatment Method`, and `Status` (only active). Also, provide the contact info for the corresponding `Tall Tree Health` location for additional assistance.
+
+7. **Practitioner's Contact Information**: Provide `Name`, `Discipline`, and `Booking Link`. Do not print their `Focus Areas`. Provide contact information in the following structured format:
+
+    - `First Name` `Last Name`:
+    - `Discipline`:
+    - `Booking Link`: (print only if available)
+
+8. **Online Booking Info**: Provide the appropriate clinic contact information from the `Tall Tree Integrated Health Centre Information` for online booking.
+
+## Tall Tree Integrated Health Service Routing Guidelines
+
+9. **Mental Health Queries**: Recommend psychologist or clinical counsellour for mental health queries, including depression, stress, anxiety, trauma, suicidal thoughts, etc.
+
+10. **Injuries and Pain**: Prioritize Physiotherapy for injuries and pain conditions unless another preference is stated.
+
+11. **Randomness in Recommendations**: Introduce randomness in practitioner recommendations for general issues to avoid bias.
+
+12. **Concussion Protocol**: Direct to the `Concussion Treatment Program` for the appropriate location for a comprehensive assessment with a physiotherapist. Do not recommend a practitioner.
+
+13. **Psychologist in Vancouver**: If a Psychologist is requested in the Vancouver location, provide only the contact and booking link for our mental health team in Cordova Bay - Upstairs location. Do not recommend an alternative practitioner.
+
+14. **Sleep issues**: Recommend only the Sleep Program intake and provide the phone number to book an appointment. Do not recommend a practitioner.
+
+15. **Longevity Program**: For longevity queries, provide the Longevity Program phone number. Do not recommend a practitioner.
+
+16. **DEXA Testing or body composition**: Inform that this service is exclusive to the Cordova Bay clinic and provide the clinic phone number and booking link. Do not recommend a practitioner.
+
+17. **For VO2 Max Testing**: Determine the patient's location preference for Vancouver or Victoria and provide the booking link for the appropriate location. If Victoria, we only do it at our Cordova Bay location.
+
+18. **Assistance and Closure**: Offer further assistance and conclude positively with a reassuring statement without being repetitive. Example: "Take care! 😊", etc.
+
+---
+
+# Patient Query
+
+```
+{message}
+```
+---
+
+# Context
+
+---
+1. **Practitioners Database**:
+
+```
+{practitioners_db}
+```
+---
+
+2. **Tall Tree Health Centre Information**:
+
+```
+{tall_tree_db}
+```
+---
+
+"""
+
+    # Template for system message with markdown formatting
+    system_message = SystemMessagePromptTemplate.from_template(
+        system_template)
+
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            system_message,
+            MessagesPlaceholder(variable_name="history"),
+            ("human", "{message}"),
+        ]
+    )
+
+    return prompt

rag_chain/retrievers_setup.py

@@ -0,0 +1,172 @@
+import os
+from functools import cache
+
+import qdrant_client
+import torch
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import EmbeddingsFilter
+from langchain_community.retrievers import QdrantSparseVectorRetriever
+from langchain_community.vectorstores import Qdrant
+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.
+
+    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"):
+        self.validate_environment_variables()
+        self.embeddings_model = embeddings_model
+        self.collection_name = collection_name
+        self.client = qdrant_client.QdrantClient(
+            url=os.getenv("QDRANT_URL"),
+            api_key=os.getenv("QDRANT_API_KEY"),
+        )
+        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}")
+
+    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)
+
+    @cache
+    def load_qdrant_collection(self):
+        """Load Qdrant collection for a given embeddings model."""
+        if self.embeddings_model == "text-embedding-ada-002":
+            self.qdrant_collection = self.set_qdrant_collection(
+                OpenAIEmbeddings(model=self.embeddings_model))
+        else:
+            raise ValueError(
+                f"Invalid embeddings model: {self.embeddings_model}. Select 'text-embedding-ada-002' from OpenAI.")
+
+        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
+
+
+class SparseRetrieverClient:
+    """Inititalize 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: 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"))
+        self.model_id = splade_model_id
+        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}")
+
+    @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
+
+    @cache
+    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
+        """
+        tokenizer, model = self.set_tokenizer_config()
+        tokens = tokenizer(text, return_tensors="pt",
+                           max_length=512, padding="max_length", truncation=True)
+        output = model(**tokens)
+        logits, attention_mask = output.logits, tokens.attention_mask
+        relu_log = torch.log(1 + torch.relu(logits))
+        weighted_log = relu_log * attention_mask.unsqueeze(-1)
+        max_val, _ = torch.max(weighted_log, dim=1)
+        vec = max_val.squeeze()
+        indices = vec.nonzero().numpy().flatten()
+        values = vec.detach().numpy()[indices]
+        return indices.tolist(), values.tolist()
+
+    def get_sparse_retriever(self):
+
+        sparse_retriever = QdrantSparseVectorRetriever(
+            client=self.client,
+            collection_name=self.collection_name,
+            sparse_vector_name=self.vector_name,
+            sparse_encoder=self.sparse_encoder,
+            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).
+
+    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 compression_retriever