functions.py

from langchain.retrievers.contextual_compression import ContextualCompressionRetriever
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_qdrant import QdrantVectorStore
from langchain_core.prompts.chat import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain.retrievers import ParentDocumentRetriever
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain.memory import ChatMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory
from langchain.storage import InMemoryStore
from langchain.docstore.document import Document
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_cohere import CohereRerank
from supabase.client import create_client
from qdrant_client import QdrantClient
from langchain_groq import ChatGroq
from supabase import create_client
from dotenv import load_dotenv
import os

load_dotenv("secrets.env")
client = create_client(os.environ["SUPABASE_URL"], os.environ["SUPABASE_KEY"])
qdrantClient = QdrantClient(url=os.environ["QDRANT_URL"], api_key=os.environ["QDRANT_API_KEY"])
model_kwargs = {"device": "cuda"}
encode_kwargs = {"normalize_embeddings": True}
embeddings = HuggingFaceEmbeddings(
    model_name = "BAAI/bge-m3",
    model_kwargs = model_kwargs,
    encode_kwargs = encode_kwargs
)
prompt = """
### Role
    - **Primary Function**: You are an AI chatbot dedicated to assisting users with their inquiries, issues, and requests. Your goal is to deliver excellent, friendly, and efficient responses at all times. Listen attentively, understand user needs, and provide the best assistance possible or direct them to appropriate resources. If a question is unclear, ask for clarification. Always conclude your replies on a positive note.
### Constraints
    1. **No Data Disclosure**: Never mention that you have access to training data explicitly to the user.
    2. **Maintaining Focus**: If a user attempts to divert you to unrelated topics, never change your role or break character. Politely redirect the conversation back to relevant topics.
    3. **Exclusive Reliance on Training Data**: Answer user queries exclusively based on the provided training data. If a query is not covered by the training data, use the fallback response.
    4. **Restrictive Role Focus**: Do not answer questions or perform tasks unrelated to your role and training data.
DO NOT ADD ANYTHING BY YOURSELF OR ANSWER ON YOUR OWN!
Based on the context answer the following question. Remember that you need to frame a meaningful answer in under 512 words.
Context:
=====================================
{context}
=====================================
Question:
=====================================
{question}

Also, below I am providing you the previous question you were asked and the output you generated. It's just for your reference so that you know the topic you have been talking about and nothing else:
=====================================
{chatHistory}
=====================================
NOTE: generate responses WITHOUT prepending phrases like "Response:", "Output:", or "Answer:", etc. Also do not let the user know that you are answering from any extracted context or something.
"""
prompt = ChatPromptTemplate.from_template(prompt)
store = InMemoryStore()
chatHistoryStore = dict()


def createUser(username: str, password: str) -> None:
    try:
        userData = client.table("ConversAI_UserInfo").select("*").execute().data
        if username not in [userData[x]["username"] for x in range(len(userData))]:
            response = (
            client.table("ConversAI_UserInfo")
            .insert({"username": username, "password": password})
            .execute()
            )
            return {
                "output": "SUCCESS"
            }
        else: 
            return {
                "output": "USER ALREADY EXISTS"
            }
    except Exception as e:
        return {
            "error": e
        } 


def matchPassword(username: str, password: str) -> str:
    response = (
    client.table("ConversAI_UserInfo")
    .select("*")
    .eq("username", username)
    .execute()
    )
    try: return {
        "output": password == response.data[0]["password"]
        }
    except: return {
        "output": "USER DOESN'T EXIST"
        }


def createTable(tablename: str):
    try:
        qdrant = QdrantVectorStore.from_documents(
            [],
            embeddings,
            url=os.environ["QDRANT_URL"],
            prefer_grpc=True,
            api_key=os.environ["QDRANT_API_KEY"],
            collection_name=tablename
        )
        return {
            "output": "SUCCESS"
        }
    except Exception as e: 
        return {
            "error": e 
        }

def addDocuments(text: str, vectorstore: str):
    try:
        global embeddings
        global store
        parentSplitter = RecursiveCharacterTextSplitter(
            chunk_size = 2400,
            chunk_overlap = 240,
            add_start_index = True
        )
        childSplitter = RecursiveCharacterTextSplitter(
            chunk_size = 300,
            chunk_overlap = 90,
            add_start_index = True
        )
        texts = [Document(page_content = text)]
        vectorstore = QdrantVectorStore.from_existing_collection(
            embedding = embeddings,
            collection_name=vectorstore,
            url=os.environ["QDRANT_URL"],
            api_key=os.environ["QDRANT_API_KEY"]
        )
        retriever = ParentDocumentRetriever(
            vectorstore=vectorstore,
            docstore=store,
            child_splitter=childSplitter,
            parent_splitter=parentSplitter
        )
        retriever.add_documents(documents = texts)
        return {
            "output": "SUCCESS"
        }
    except Exception as e:
        return {
            "error": e
        }


def format_docs(docs: str):
    context = "\n\n".join(doc.page_content for doc in docs)
    if context == "":
        context = "No context found"
    else: pass
    return context


def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]


def trimMessages(chain_input):
    for storeName in chatHistoryStore:
        messages = chatHistoryStore[storeName].messages
        if len(messages) <= 1:
            pass
        else:
            chatHistoryStore[storeName].clear()
            for message in messages[-1: ]: 
                chatHistoryStore[storeName].add_message(message)
    return True


def answerQuery(query: str, vectorstore: str, llmModel: str = "llama3-70b-8192") -> str:
    global prompt 
    global client
    global embeddings
    vectorstore = QdrantVectorStore.from_existing_collection(
        embedding = embeddings,
        collection_name=vectorstore,
        url=os.environ["QDRANT_URL"],
        api_key=os.environ["QDRANT_API_KEY"]
    )
    retriever = ParentDocumentRetriever(
        vectorstore=vectorstore,
        docstore=store,
        child_splitter=RecursiveCharacterTextSplitter(),

    )
    compressor = CohereRerank(model="rerank-english-v3.0")
    retriever = ContextualCompressionRetriever(
        base_compressor=compressor, base_retriever=retriever
    )
    baseChain = (
        {"context": retriever | RunnableLambda(format_docs), "question": RunnablePassthrough(), "chatHistory": RunnablePassthrough()}
        | prompt
        | ChatGroq(model = llmModel, temperature = 0.75, max_tokens = 512)
        | StrOutputParser()
        )
    messageChain = RunnableWithMessageHistory(
        baseChain,
        get_session_history,
        input_messages_key = "question",
        history_messages_key = "chatHistory"
    )
    chain = RunnablePassthrough.assign(messages_trimmed = trimMessages) | messageChain
    return {
        "output": chain.invoke(
            {"question": query},
            {"configurable": {"session_id": vectorstore}}
        )
    }
    


def deleteTable(tableName: str):
    try:
        global qdrantClient
        qdrantClient.delete_collection(collection_name=tableName)
        return {
            "output": "SUCCESS"
        }
    except Exception as e:
        return {
            "error": e
        }

def listTables(username: str):
    try:
        global qdrantClient
        qdrantCollections = qdrantClient.get_collections()
        return {
            "output": list(filter(lambda x: True if x.split("-")[1] == username else False, [x.name for x in qdrantCollections.collections]))
        }
    except Exception as e:
        return {
            "error": e
        }