get_similarity.py

from src.pgdb.knowledge.similarity import VectorStore_FAISS
from src.config.prompts import PROMPT_QUERY_EXTEND,PROMPT_QA_EXTEND_QUESTION
from .rerank import BgeRerank


from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers.list import ListOutputParser


class GetSimilarity:
    def __init__(self, _question, _faiss_db: VectorStore_FAISS):
        self.question = _question
        self.faiss_db = _faiss_db
        self.similarity_docs = self.faiss_db.get_text_similarity(self.question)
        self.similarity_doc_txt = self.faiss_db.join_document(self.similarity_docs)
        self.rerank_docs = ""

    def get_rerank(self,reranker:BgeRerank ,top_k = 5):
        rerank_docs = reranker.compress_documents(self.similarity_docs,self.question)
        d_list = []
        for d in rerank_docs[:top_k]:
            d_list.append(d)
        self.rerank_docs = rerank_docs[:top_k]
        return self.faiss_db.join_document(d_list)

    def get_similarity_doc(self):
        return self.similarity_doc_txt
    
    def get_similarity_docs(self):
        return self.similarity_docs
    
    def get_rerank_docs(self):
        return self.rerank_docs


class GetSimilarityWithExt:
    def __init__(self, _question, _faiss_db: VectorStore_FAISS):
        self.question = _question
        self.faiss_db = _faiss_db
        self.similarity_docs = self.get_text_similarity_with_ext()
        self.similarity_doc_txt = self.faiss_db.join_document(self.similarity_docs)
        self.rerank_docs = ""

    def get_rerank(self, reranker: BgeRerank, top_k=5):
        question = '\n'.join(self.question)
        print(question)
        rerank_docs = reranker.compress_documents(self.similarity_docs, question)
        d_list = []
        for d in rerank_docs[:top_k]:
            d_list.append(d)
        self.rerank_docs = rerank_docs[:top_k]
        return self.faiss_db.join_document(d_list)

    def get_similarity_doc(self):
        return self.similarity_doc_txt

    def get_similarity_docs(self):
        return self.similarity_docs

    def get_rerank_docs(self):
        return self.rerank_docs

    def get_text_similarity_with_ext(self):
        similarity_docs = []
        for q in self.question:
            similarity_doc = self.faiss_db.get_text_similarity(q)
            similarity_docs.extend(similarity_doc)
        content_set = set()
        unique_documents = []
        for doc in similarity_docs:
            content = hash(doc.page_content)
            if content not in content_set:
                unique_documents.append(doc)
                content_set.add(content)
        return unique_documents

class QAExt:
    llm = None

    def __init__(self, llm) -> None:
        self.llm = llm
        prompt = PromptTemplate.from_template(PROMPT_QUERY_EXTEND)
        # parser = ListOutputParser()
        self.query_extend = prompt | llm 

    def extend_query(self, question, messages=None):
        """
            question: str
            messages: list of tuple (str,str)
                eg:
                [
                    ("Q1","A1"),
                    ("Q2","A2"),
                    ...
                ]
            
        """
        if not messages:
            messages = []
        history = ""
        for msg in messages:
            history += f"Q: {msg[0]}\nA: {msg[1]}\n"
        return self.query_extend.invoke(input={"histories":messages, "query":question})

class ChatExtend:
    def __init__(self, llm) -> None:
        self.llm = llm
        prompt = PromptTemplate.from_template(PROMPT_QA_EXTEND_QUESTION)
        self.query_extend = prompt | llm

    def new_questions(self, messages):
        history = ""
        for msg in messages:
            history += f"Q: {msg[0]}\nA: {msg[1]}\n"
        return self.query_extend.invoke(input={"histories":history})