get_similarity.py

from src.pgdb.knowledge.similarity import VectorStore_FAISS
from src.config.prompts import PROMPT_QUERY_EXTEND,PROMPT_QA_EXTEND_QUESTION
from src.server.rerank import BgeRerank,reciprocal_rank_fusion
import time
from typing import List
from langchain_core.documents import Document

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, _location=None):
        self.question = _question
        self.faiss_db = _faiss_db
        self.location = _location
        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 join_document(self, docs: List[Document]) -> str:
        if len(docs) == 0:
            return "[]"
        result = "["
        d1 = docs[0]
        page_number = "0"
        if "page_number" in d1.metadata:
            page_number = str(d1.metadata["page_number"])
        hash_str = ""
        if "hash" in d1.metadata:
            hash_str = d1.metadata["hash"]
        result += ("{\"page_content\": \"" + d1.page_content + "\",\"from_file\":\"" + d1.metadata["filename"] + "\",\"hash\":\"" + hash_str + "\",\"page_number\":\"" + page_number + "\"}")
        docs = docs[1:]
        for doc in docs:
            page_number = "0"
            if "page_number" in doc.metadata:
                page_number = str(doc.metadata["page_number"])
            hash_str = ""
            if "hash" in doc.metadata:
                hash_str = doc.metadata["hash"]
            result += (",{\"page_content\": \"" + doc.page_content + "\",\"from_file\":\"" + doc.metadata["filename"] + "\",\"hash\":\"" + hash_str + "\",\"page_number\":\"" + page_number + "\"}")
        result += "]"
        return result

    def get_rerank_with_doc(self, reranker: BgeRerank,split_docs_list:list):
        # top_k = self.get_doc_nums(len(split_docs_list))
        question = '\n'.join(self.question)
        print(question)
        # rerank_docs1_hash = []
        # rerank_docs2_hash = []
        # m = {}
        result = []
        for split_doc in split_docs_list:
            start = time.time()
            rerank_docs1 = reranker.compress_documents(split_doc, question)
            result.extend(rerank_docs1[:3])
            end = time.time()
            print('重排1 time: %s Seconds' % (end - start))
            # for doc in rerank_docs1:
            #     m[hash(doc.page_content)] = doc
            #     rerank_docs1_hash.append(hash(doc.page_content))
            # result.append((60, rerank_docs1_hash))
        start = time.time()
        rerank_docs2 = reranker.compress_documents(self.similarity_docs, question)
        result.extend(rerank_docs2[:3])
        end = time.time()
        print('重排2 time: %s Seconds' % (end - start))
        print(len(result))
        content_set = set()
        unique_documents = []
        for doc in result:
            content = hash(doc.page_content)
            if content not in content_set:
                unique_documents.append(doc)
                content_set.add(content)
        # for doc in rerank_docs2:
        #     m[hash(doc.page_content)] = doc
        #     rerank_docs2_hash.append(hash(doc.page_content))
        #
        # result.append((55,rerank_docs2_hash))
        # print(len(rerank_docs1_hash))
        # print(len(rerank_docs2_hash))
        # start = time.time()
        # rrf_doc = reciprocal_rank_fusion(result)
        # end = time.time()
        # print('混排 time: %s Seconds' % (end - start))
        # print("混排文档数量：", len(rrf_doc))
        # d_list = []
        # for key in rrf_doc:
        #     d_list.append(m[key])
        # print("返回文档数量：",top_k)
        print(len(unique_documents))
        self.rerank_docs = unique_documents
        return self.join_document(unique_documents)

    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:
            if self.location is not None:
                if any(substring in doc.page_content for substring in self.location) or any(substring in doc.metadata["filename"] for substring in self.location) :
                    content = hash(doc.page_content)
                    if content not in content_set:
                        unique_documents.append(doc)
                        content_set.add(content)
        print(len(unique_documents))
        return unique_documents

    def get_doc_nums(self,num :int)->int:
        num = num*3
        if num<5:
            return 5
        elif num>30:
            return 30
        else:
            return num

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})

    def extend_query_with_str(self, question, messages):
        """
            question: str
            messages: list of tuple (str,str)
                eg:
                [
                    ("Q1","A1"),
                    ("Q2","A2"),
                    ...
                ]

        """
        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})