emb.py

import os
from langchain.document_loaders import PyPDFLoader, DirectoryLoader, PDFMinerLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import SentenceTransformerEmbeddings
from langchain.vectorstores import Chroma
import configparser
from tqdm import tqdm
from langchain.vectorstores import Pinecone
from langchain.schema import Document
import pinecone
from dotenv import load_dotenv
from llm import LLMManager

class EmbeddingsManager:

    def __init__(self,settings, emb="hkunlp/instructor-large"):

        #Loading env variables
        load_dotenv()

        #Loading config file
        self.config=configparser.ConfigParser()
        self.config.read("config.ini")

        #Loading settingManager
        self.set=settings

        #Loading default parameters for search
        self.search_method=self.set.search_method
        self.n_doc_return=self.set.n_doc_return
        self.ai_assisted_search=self.config.getboolean('RAG','default_ai_assisted_search')
        self.available_search_methods=self.set.available_search_methods
        self.text_split_size=self.config.getint('RAG','default_text_split_size')
        self.text_overlap=self.config.getint('RAG','default_text_overlap')

        #Loading available Vector Stores
        self.vector_stores=self.get_vector_list()
        self.vector_stores_map=self.get_vector_map_list()

        #Selecting the embeddings model
        self.embedding_model_name=emb

        #Initing
        current_dir = os.path.dirname(__file__)
        data_dir = os.path.join(current_dir, "data")
        os.environ['TRANSFORMERS_CACHE'] = data_dir
        PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY')
        PINECONE_API_ENV = os.environ.get('PINECONE_API_ENV')
        self.embeddings_model = SentenceTransformerEmbeddings(model_name=self.embedding_model_name, cache_folder=data_dir)
        pinecone.init(api_key=PINECONE_API_KEY,environment=PINECONE_API_ENV)

    #This function used to get the list of emb
    def get_emb_list(self):
        """Returns a list of the available Embedding models"""
        emb_map_section = 'EMB'
        if emb_map_section in self.config:
            return [self.config.get(emb_map_section, emb) for emb in self.config[emb_map_section]]
        else:
            return []
        
    #This function used to get the list of available VectorStores
    def get_vector_list(self):
        """Returns a list of the available Vector Stores"""
        section = 'Vector_Stores'
        if section in self.config:
            return [self.config.get(section, vector) for vector in self.config[section]]
        else:
            return []
        
    #This function used to get the map of available VectorStores
    def get_vector_map_list(self):
        """Returns a list of the available Vector Stores"""
        section = 'Vector_Stores_Map'
        if section in self.config:
            return [self.config.get(section, vector) for vector in self.config[section]]
        else:
            return []

    #This function is used to get the relevant context
    def get_context(self,index, query, history):
        """Returns the relevant context for the LLM"""

        docsearch = Pinecone.from_existing_index(index, self.embeddings_model)

        if self.set.ai_assisted_search:
            prompt=self.set.default_ai_search_prompt
            prompt=prompt.format(question=query,history=history)
            print(prompt)
            llm=LLMManager(self.set)
            queryterms=llm.get_query_terms(prompt)
            query=queryterms+"\n"+query
            
            #print("new query input={new_query}".format(new_query=query))
        

        if self.set.search_method=="MMR":
            return docsearch.max_marginal_relevance_search(query, k=self.set.n_doc_return,fetch_metadata=True)
        
        elif self.set.search_method=="Similarity":
            return docsearch.similarity_search(query, k=self.set.n_doc_return,fetch_metadata=True)
        
        else:
            return docsearch.max_marginal_relevance_search(query, k=self.set.n_doc_return,fetch_metadata=True)
        

        #This function is used to get the relevant context
    def get_context_search(self,index, query):
        """Returns the relevant context for the LLM"""

        docsearch = Pinecone.from_existing_index(index, self.embeddings_model)

        if self.set.search_method=="MMR":
            return docsearch.max_marginal_relevance_search(query, k=2,fetch_metadata=True)
        
        elif self.set.search_method=="Similarity":
            return docsearch.similarity_search(query, k=2,fetch_metadata=True)
        
        else:
            return docsearch.max_marginal_relevance_search(query, k=self.n_doc_return,fetch_metadata=True)
    
    #This function is used to get the relevant context formatted
    def get_formatted_context(self,index, query,history):
        """Returns the relevant context for the LLM formatted"""

        formatted=""
        docs=self.get_context(index, query,history)
        for doc in docs:
            formatted+="DOCUMENT NAME={doc_name}\nDOCUMENT CONTENT={doc_content}\n\n".format(doc_name=doc.metadata["source"],doc_content=doc.page_content)
        return formatted



    #This function is used to add documents to an existing vector store
    def generate_vector_store(self, index):
        """Adds a document to the vector store on Pinecone."""

        documents = []
        for root, dirs, files in os.walk("docs"):
            for file in files:
                if file.endswith(".pdf"):
                    print("Uploading "+file.replace(".pdf",""))
                    documents.clear()
                    loader = PDFMinerLoader(os.path.join(root, file))
                    documents.extend(loader.load())
                    text_splitter = RecursiveCharacterTextSplitter(chunk_size=self.text_split_size, chunk_overlap=self.text_overlap)
                    texts = text_splitter.split_documents(documents)
                    docsearch = Pinecone.from_documents(texts, embedding=self.embeddings_model, index_name=index)
                    os.remove(os.path.join(root, file))

        return "Ok"
    

    # Example Usage:
if __name__ == "__main__":

    """This is an example of how to add document to the vectorstore on Pinecone"""
    from settings import SettingManager
    set= SettingManager()
    emb_manager = EmbeddingsManager(set,emb="hkunlp/instructor-large")
    print(emb_manager.generate_vector_store("prohelper"))


    #"""This is an example of how to retrive context and display all values retrived"""
    #emb_manager = EmbeddingsManager()
    #docs=emb_manager.get_context(index="prohelper",query="Could you explain to me what is esrs?")
    #for i in docs:
    #    print("---------------------------------------------------------")
    #    print(i.metadata["Doc"])
    #    print("           ")
    #    print(i.page_content)
    #llm_manager.selectLLM("Mixtral 7B")