Upload 4 files

rag/helpers.py

@@ -0,0 +1,57 @@
+import logging
+import os
+import random
+from datetime import datetime
+from functools import lru_cache
+from typing import Sequence
+from zoneinfo import ZoneInfo
+
+import langsmith
+from langchain_core.documents import Document
+from langchain_community.document_transformers import LongContextReorder
+from langchain.retrievers.document_compressors import FlashrankRerank
+
+logging.basicConfig(level=logging.ERROR)
+
+
+class DocumentFormatter:
+    def __init__(self, prefix: str):
+        self.prefix = prefix
+
+    def __call__(self, docs: list[Document]) -> str:
+        return "\n---\n".join(
+            [
+                f"- {self.prefix} {i+1}:\n\n\t" + d.page_content
+                for i, d in enumerate(docs)
+            ]
+        )
+
+
+def get_datetime() -> str:
+    return datetime.now(ZoneInfo("America/Vancouver")).strftime("%A, %Y-%b-%d %H:%M:%S")
+
+
+def reorder_documents(docs: list[Document]) -> Sequence[Document]:
+    return LongContextReorder().transform_documents(docs)
+
+
+def randomize_documents(documents: list[Document]) -> list[Document]:
+    random.shuffle(documents)
+    return documents
+
+
+def create_langsmith_client():
+    os.environ["LANGCHAIN_TRACING_V2"] = "true"
+    os.environ["LANGCHAIN_PROJECT"] = "admin-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()
+
+
+@lru_cache(maxsize=1)
+def get_reranker(
+    top_n: int = 3, model: str = "ms-marco-MiniLM-L-12-v2"
+) -> FlashrankRerank:
+    return FlashrankRerank(top_n=top_n, model=model)

rag/prompt_template.py

@@ -11,7 +11,7 @@ def generate_prompt_template():
 
 ---
 
-Your name is Ella (Empathetic, Logical, Liaison, Accessible). You are a helpful Virtual Assistant at Tall Tree Health in British Columbia, Canada. Based on the patient's symptoms/needs, connect them with the right 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:
+Your name is ELLA (Empathetic, Logical, Liaison, Accessible). You are a helpful AI Assistant at Tall Tree Health in British Columbia, Canada. Based on the patient's symptoms/needs, connect them with the right 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:
 
 ---
 
@@ -58,7 +58,7 @@ Your name is Ella (Empathetic, Logical, Liaison, Accessible). You are a helpful
 # Patient Query
 
 ```
-{message}
+{user_query}
 ```
 ---
 
@@ -81,14 +81,14 @@ Your name is Ella (Empathetic, Logical, Liaison, Accessible). You are a helpful
 
 """
 
-    # Template for system message with markdown formatting
+    # Template for the system message with markdown formatting
     system_message = SystemMessagePromptTemplate.from_template(system_template)
 
     prompt = ChatPromptTemplate.from_messages(
         [
             system_message,
             MessagesPlaceholder(variable_name="history"),
-            ("human", "{message}"),
+            ("human", "{user_query}"),
         ]
     )
 

rag/retrievers.py

@@ -1,5 +1,6 @@
 import os
-from typing import List, Literal
+from functools import lru_cache
+from typing import Literal
 
 from langchain_core.vectorstores import VectorStoreRetriever
 from langchain_openai import OpenAIEmbeddings
@@ -9,6 +10,8 @@ os.environ["GRPC_VERBOSITY"] = "NONE"
 
 
 class RetrieversConfig:
+    REQUIRED_ENV_VARS = ["QDRANT_API_KEY", "QDRANT_URL", "OPENAI_API_KEY"]
+
     def __init__(
         self,
         dense_model_name: Literal["text-embedding-3-small"] = "text-embedding-3-small",
@@ -16,24 +19,35 @@ class RetrieversConfig:
             "prithivida/Splade_PP_en_v1"
         ] = "prithivida/Splade_PP_en_v1",
     ):
-        self.required_env_vars = ["QDRANT_API_KEY", "QDRANT_URL", "OPENAI_API_KEY"]
-        self._validate_environment(self.required_env_vars)
+        self._validate_environment()
         self.qdrant_url = os.getenv("QDRANT_URL")
         self.qdrant_api_key = os.getenv("QDRANT_API_KEY")
-        self.dense_embeddings = OpenAIEmbeddings(model=dense_model_name)
-        self.sparse_embeddings = FastEmbedSparse(
-            model_name=sparse_model_name,
-        )
+        self.dense_model_name = dense_model_name
+        self.sparse_model_name = sparse_model_name
 
-    def _validate_environment(self, required_env_vars: List[str]):
+    @staticmethod
+    def _validate_environment():
         missing_vars = [
-            var for var in required_env_vars if not os.getenv(var, "").strip()
+            var
+            for var in RetrieversConfig.REQUIRED_ENV_VARS
+            if not os.getenv(var, "").strip()
         ]
         if missing_vars:
             raise EnvironmentError(
                 f"Missing or empty environment variable(s): {', '.join(missing_vars)}"
             )
 
+    @property
+    @lru_cache(maxsize=2)
+    def dense_embeddings(self):
+        return OpenAIEmbeddings(model=self.dense_model_name)
+
+    @property
+    @lru_cache(maxsize=2)
+    def sparse_embeddings(self):
+        return FastEmbedSparse(model_name=self.sparse_model_name)
+
+    @lru_cache(maxsize=8)
     def get_qdrant_retriever(
         self,
         collection_name: str,
@@ -55,14 +69,6 @@ class RetrieversConfig:
 
         return qdrantdb.as_retriever(search_kwargs={"k": k})
 
-    def get_documents_retriever(self, k: int = 5) -> VectorStoreRetriever:
-        return self.get_qdrant_retriever(
-            collection_name="docs_hybrid_db",
-            dense_vector_name="docs_dense_vectors",
-            sparse_vector_name="docs_sparse_vectors",
-            k=k,
-        )
-
     def get_practitioners_retriever(self, k: int = 5) -> VectorStoreRetriever:
         return self.get_qdrant_retriever(
             collection_name="practitioners_hybrid_db",
@@ -70,3 +76,11 @@ class RetrieversConfig:
             sparse_vector_name="practitioners_sparse_vectors",
             k=k,
         )
+
+    def get_documents_retriever(self, k: int = 5) -> VectorStoreRetriever:
+        return self.get_qdrant_retriever(
+            collection_name="docs_hybrid_db",
+            dense_vector_name="docs_dense_vectors",
+            sparse_vector_name="docs_sparse_vectors",
+            k=k,
+        )

rag/runnable_and_memory.py

@@ -0,0 +1,107 @@
+import logging
+from operator import itemgetter
+
+from langchain.memory import ConversationBufferWindowMemory
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import Runnable, RunnableLambda
+from langchain_openai import ChatOpenAI
+
+from rag.retrievers import RetrieversConfig
+
+from .helpers import (
+    DocumentFormatter,
+    create_langsmith_client,
+    get_datetime,
+    get_reranker,
+)
+from .prompt_template import generate_prompt_template
+
+logging.basicConfig(level=logging.ERROR)
+
+
+def retrievers_setup(retrievers_config, reranker: bool = False) -> tuple:
+    """Set up retrievers with re-ranking
+    Args:
+        retrievers_config (_type_):
+        reranker (bool, optional): Defaults to False.
+
+    Returns:
+        tuple: Retrievers
+    """
+    # Practitioners
+    practitioners_retriever = retrievers_config.get_practitioners_retriever(k=10)
+    # Tall Tree documents
+    documents_retriever = retrievers_config.get_documents_retriever(k=10)
+
+    # Re-ranking (optional): Improves quality and serves as a filter
+    if reranker:
+        practitioners_retriever_reranker = ContextualCompressionRetriever(
+            base_compressor=get_reranker(top_n=10),
+            base_retriever=practitioners_retriever,
+        )
+        documents_retriever_reranker = ContextualCompressionRetriever(
+            base_compressor=get_reranker(top_n=8),
+            base_retriever=documents_retriever,
+        )
+
+        return practitioners_retriever_reranker, documents_retriever_reranker
+
+    else:
+        return practitioners_retriever, documents_retriever
+
+
+# Set retrievers as global variables (I see better loading time from Streamlit this way)
+practitioners_retriever, documents_retriever = retrievers_setup(
+    retrievers_config=RetrieversConfig(), reranker=True
+)
+
+
+# Set up runnable and chat memory
+def get_runnable_and_memory(
+    model: str = "gpt-4o-mini", temperature: float = 0.1
+) -> tuple[Runnable, ConversationBufferWindowMemory]:
+    """Set up runnable and chat memory
+
+    Args:
+        model_name (str, optional): LLM model. Defaults to "gpt-4o-mini".
+        temperature (float, optional): Model temperature. Defaults to 0.1.
+
+    Returns:
+        Runnable, Memory: Runnable and Memory
+    """
+
+    # Set up Langsmith to trace the runnable
+    create_langsmith_client()
+
+    # LLM and prompt template
+    llm = ChatOpenAI(
+        model=model,
+        temperature=temperature,
+    )
+
+    prompt = generate_prompt_template()
+
+    # Set conversation history window memory. It only uses the last k interactions
+    memory = ConversationBufferWindowMemory(
+        memory_key="history",
+        return_messages=True,
+        k=6,
+    )
+
+    # Set up runnable using LCEL
+    setup = {
+        "practitioners_db": itemgetter("user_query")
+        | practitioners_retriever
+        | DocumentFormatter("Practitioner #"),
+        "tall_tree_db": itemgetter("user_query")
+        | documents_retriever
+        | DocumentFormatter("No."),
+        "timestamp": lambda _: get_datetime(),
+        "history": RunnableLambda(memory.load_memory_variables) | itemgetter("history"),
+        "user_query": itemgetter("user_query"),
+    }
+
+    runnable = setup | prompt | llm | StrOutputParser()
+
+    return runnable, memory