app/utils/chat_rag.py

#list of librarys for requirement.txt
import os
import re
import hashlib
import asyncio

from langchain_community.document_loaders import PyPDFLoader

# Import embeddings module from langchain_community for vector representations of text
from langchain_community.embeddings import HuggingFaceEmbeddings

# Import text splitter for handling large texts
from langchain.text_splitter import CharacterTextSplitter

# Import vector store for database operations
from langchain_community.vectorstores import Chroma

# for loading of llama gguf model
from langchain_community.llms import LlamaCpp

from langchain.chains.router.llm_router import LLMRouterChain, RouterOutputParser
from langchain.chains.router.multi_prompt_prompt import MULTI_PROMPT_ROUTER_TEMPLATE

from langchain.chains.router import MultiPromptChain
from langchain.chains import ConversationChain
from langchain.chains.llm import LLMChain
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationBufferMemory, VectorStoreRetrieverMemory
from langchain.chains import ConversationalRetrievalChain

from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler



def sanitize_collection_name(email):
    # Replace invalid characters with an underscore
    sanitized = re.sub(r'[^a-zA-Z0-9_-]', '_', email)
    # Ensure the name is within the length limits
    if len(sanitized) > 63:
        # Hashing the name to ensure uniqueness and length constraint
        hash_suffix = hashlib.sha256(email.encode()).hexdigest()[:8]
        sanitized = sanitized[:55] + "_" + hash_suffix
    # Ensure it starts and ends with an alphanumeric character
    if not re.match(r'^[a-zA-Z0-9].*[a-zA-Z0-9]$', sanitized):
        sanitized = "a" + sanitized + "1"
    return sanitized


# Modify vectordb initialization to be dynamic based on user_id
async def get_vectordb_for_user(user_collection_name):
    # Get Chromadb location
    CHROMADB_LOC = os.getenv('CHROMADB_LOC')

    vectordb = Chroma(
        collection_name=user_collection_name,
        embedding_function=HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2'),
        persist_directory=f"{CHROMADB_LOC}" # Optional: Separate directory for each user's data
    )
    return vectordb

vectordb_cache = {}

async def get_vectordb_for_user_cached(user_collection_name):
    if user_collection_name not in vectordb_cache:
        vectordb_cache[user_collection_name] = await get_vectordb_for_user(user_collection_name)
    return vectordb_cache[user_collection_name]


def pdf_to_vec(filename, user_collection_name):
    
    # Get Chromadb location
    CHROMADB_LOC = os.getenv('CHROMADB_LOC')

    document = []
    loader = PyPDFLoader(filename)
    document.extend(loader.load()) #which library is this from?

    # Initialize HuggingFaceEmbeddings with the 'sentence-transformers/all-MiniLM-L6-v2' model for generating text embeddings
    embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')

    # Initialize a CharacterTextSplitter to split the loaded documents into smaller chunks
    document_splitter = CharacterTextSplitter(separator='\n', chunk_size=500, chunk_overlap=100)

    # Use the splitter to divide the 'document' content into manageable chunks
    document_chunks = document_splitter.split_documents(document) #which library is this from?

    # Create a Chroma vector database from the document chunks with the specified embeddings, and set a directory for persistence
    vectordb = Chroma.from_documents(document_chunks, embedding=embeddings, collection_name=user_collection_name, persist_directory=CHROMADB_LOC) ## change to GUI path

    # Persist the created vector database to disk in the specified directory
    vectordb.persist() #this is mandatory?

    return(vectordb)
    #return collection  # Return the collection as the asset


# Assuming LlamaModelSingleton is updated to support async instantiation
class LlamaModelSingleton:
    _instance = None

    @classmethod
    async def get_instance(cls):
        if cls._instance is None:
            cls._instance = cls._load_llm()  # Assuming _load_llm is synchronous, if not, use an executor
        return cls._instance

    @staticmethod
    def _load_llm():
        print('Loading LLM model...')
        model_path = os.getenv("MODEL_PATH")
        llm = LlamaCpp(
            model_path=model_path,
            n_gpu_layers=-1,
            n_batch=512,
            temperature=0.1,
            top_p=1,
            max_tokens=2000,
        )
        print(f'Model loaded from {model_path}')
        return llm

async def load_llm():
    return await LlamaModelSingleton.get_instance()




#step 5, to instantiate once to create default_chain,router_chain,destination_chains into chain and set vectordb. so will not re-create per prompt
async def default_chain(llm, user_collection_name):
    # Get Chromadb location
    CHROMADB_LOC = os.getenv('CHROMADB_LOC')

    vectordb = await get_vectordb_for_user_cached(user_collection_name)  # Use the dynamic vectordb based on user_id
    sum_template = """
    As a machine learning education specialist, our expertise is pivotal in deepening the comprehension of complex machine learning concepts for both educators and students.

    our role entails:

    Providing Detailed Explanations: Deliver comprehensive answers to these questions, elucidating the underlying technical principles.
    Assisting in Exam Preparation: Support educators in formulating sophisticated exam and quiz questions, including MCQs, accompanied by thorough explanations.
    Summarizing Course Material: Distill key information from course materials, articulating complex ideas within the context of advanced machine learning practices.

    Objective: to summarize and explain the key points.
    Here the question:
    {input}"""

    mcq_template = """
    As a machine learning education specialist, our expertise is pivotal in deepening the comprehension of complex machine learning concepts for both educators and students.

    our role entails:
    Crafting Insightful Questions: Develop thought-provoking questions that explore the intricacies of machine learning topics.
    Generating MCQs: Create MCQs for each machine learning topic, comprising a question, four choices (A-D), and the correct answer, along with a rationale explaining the answer.

    Objective: to create multiple choice question in this format
    [question:
    options A:
    options B:
    options C:
    options D:
    correct_answer:
    explanation:]

    Here the question:
    {input}"""

    prompt_infos = [
        {
            "name": "SUMMARIZE",
            "description": "Good for summarizing and explaination ",
            "prompt_template": sum_template,
        },
        {
            "name": "MCQ",
            "description": "Good for creating multiple choices questions",
            "prompt_template": mcq_template,
        },
    ]

    destination_chains = {}

    for p_info in prompt_infos:
        name = p_info["name"]
        prompt_template = p_info["prompt_template"]
        #vectordb=p_info["vector"]
        prompt = PromptTemplate(template=prompt_template, input_variables=["input"])
        embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2') #new
        vectordb= Chroma(persist_directory = CHROMADB_LOC, embedding_function = embeddings) #new
        retriever = vectordb.as_retriever()#new
        memory = VectorStoreRetrieverMemory(retriever=retriever) #new

        chain = LLMChain(llm=llm, prompt=prompt, verbose=True, memory=memory) #new memory=memory
        destination_chains[name] = chain

    #default_chain = ConversationChain(llm=llm, output_key="text")
    #memory = ConversationBufferMemory(memory_key='chat_history', return_messages=True)

    default_chain = ConversationalRetrievalChain.from_llm(llm=llm,
                                                retriever=vectordb.as_retriever(search_kwargs={'k': 3}),
                                                verbose=True, output_key="text" )

    destinations = [f"{p['name']}: {p['description']}" for p in prompt_infos]
    destinations_str = "\n".join(destinations)
    router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(destinations=destinations_str)
    router_prompt = PromptTemplate(
        template=router_template,
        input_variables=["input"],
        output_parser=RouterOutputParser(),
    )
    router_chain = LLMRouterChain.from_llm(llm, router_prompt)

    return default_chain,router_chain,destination_chains

# Adjust llm_infer to accept user_id and use it for user-specific processing
async def llm_infer(user_collection_name, prompt):
    
    llm = await load_llm()  # load_llm is singleton for entire system
    
    vectordb = await get_vectordb_for_user_cached(user_collection_name) # Vector collection for each us. 
    
    default_chain, router_chain, destination_chains = await get_or_create_chain(user_collection_name, llm)  # Now user-specific

    chain = MultiPromptChain(
        router_chain=router_chain,
        destination_chains=destination_chains,
        default_chain=default_chain,
        #memory=ConversationBufferMemory(k=2), # memory_key='chat_history', return_messages=True
        verbose=True,
    )
    response = chain.run(prompt)

    return response

# Assuming a simplified caching mechanism for demonstration
chain_cache = {}

async def get_or_create_chain(user_collection_name, llm):
    if 'default_chain' in chain_cache and 'router_chain' in chain_cache:
        default_chain = chain_cache['default_chain']
        router_chain = chain_cache['router_chain']
        destination_chains = chain_cache['destination_chains']
    else:
        vectordb = await get_vectordb_for_user_cached(user_collection_name)  # User-specific vector database
        sum_template = """
        As a machine learning education specialist, our expertise is pivotal in deepening the comprehension of complex machine learning concepts for both educators and students.

        our role entails:

        Providing Detailed Explanations: Deliver comprehensive answers to these questions, elucidating the underlying technical principles.
        Assisting in Exam Preparation: Support educators in formulating sophisticated exam and quiz questions, including MCQs, accompanied by thorough explanations.
        Summarizing Course Material: Distill key information from course materials, articulating complex ideas within the context of advanced machine learning practices.

        Objective: to summarize and explain the key points.
        Here the question:
        {input}"""

        mcq_template = """
        As a machine learning education specialist, our expertise is pivotal in deepening the comprehension of complex machine learning concepts for both educators and students.

        our role entails:
        Crafting Insightful Questions: Develop thought-provoking questions that explore the intricacies of machine learning topics.
        Generating MCQs: Create MCQs for each machine learning topic, comprising a question, four choices (A-D), and the correct answer, along with a rationale explaining the answer.

        Objective: to create multiple choice question in this format
        [question:
        options A:
        options B:
        options C:
        options D:
        correct_answer:
        explanation:]

        Here the question:
        {input}"""

        prompt_infos = [
            {
                "name": "SUMMARIZE",
                "description": "Good for summarizing and explaination ",
                "prompt_template": sum_template,
            },
            {
                "name": "MCQ",
                "description": "Good for creating multiple choices questions",
                "prompt_template": mcq_template,
            },
        ]

        destination_chains = {}

        for p_info in prompt_infos:
            name = p_info["name"]
            prompt_template = p_info["prompt_template"]
            prompt = PromptTemplate(template=prompt_template, input_variables=["input"])
            chain = LLMChain(llm=llm, prompt=prompt)
            destination_chains[name] = chain
        #default_chain = ConversationChain(llm=llm, output_key="text")
        #memory = ConversationBufferMemory(memory_key='chat_history', return_messages=True)

        default_chain = ConversationalRetrievalChain.from_llm(llm=llm,
                                                    retriever=vectordb.as_retriever(search_kwargs={'k': 3}),
                                                    verbose=True, output_key="text" )

        destinations = [f"{p['name']}: {p['description']}" for p in prompt_infos]
        destinations_str = "\n".join(destinations)
        router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(destinations=destinations_str)
        router_prompt = PromptTemplate(
            template=router_template,
            input_variables=["input"],
            output_parser=RouterOutputParser(),
        )
        router_chain = LLMRouterChain.from_llm(llm, router_prompt)
#
        chain_cache['default_chain'] = default_chain
        chain_cache['router_chain'] = router_chain
        chain_cache['destination_chains'] = destination_chains
    
    # Here we can adapt the chains if needed based on the user_id, for example, by adjusting the vectordb retriever
    # This is where user-specific adaptations occur

    return default_chain, router_chain, destination_chains