get_similarity.py

from src.pgdb.knowledge.similarity import VectorStore_FAISS
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


DEFAULT_PROMPT = """作为一个向量检索助手，你的任务是结合历史记录，从不同角度，为“原问题”生成个不同版本的“检索词”，从而提高向量检索的语义丰富度，提高向量检索的精度。生成的问题要求指向对象清晰明确，并与“原问题语言相同”。例如：
历史记录: 
'''
'''
原问题: 介绍下剧情。
检索词: ["介绍下故事的背景。","故事的主题是什么？","介绍下故事的主要人物。"]
----------------
历史记录: 
'''
Q: 对话背景。
A: 当前对话是关于 Nginx 的介绍和使用等。
'''
原问题: 怎么下载
检索词: ["Nginx 如何下载？","下载 Nginx 需要什么条件？","有哪些渠道可以下载 Nginx？"]
----------------
历史记录: 
'''
Q: 对话背景。
A: 当前对话是关于 Nginx 的介绍和使用等。
Q: 报错 "no connection"
A: 报错"no connection"可能是因为……
'''
原问题: 怎么解决
检索词: ["Nginx报错"no connection"如何解决？","造成'no connection'报错的原因。","Nginx提示'no connection'，要怎么办？"]
----------------
历史记录: 
'''
Q: 护产假多少天?
A: 护产假的天数根据员工所在的城市而定。请提供您所在的城市，以便我回答您的问题。
'''
原问题: 沈阳
检索词: ["沈阳的护产假多少天？","沈阳的护产假政策。","沈阳的护产假标准。"]
----------------
历史记录: 
'''
Q: 作者是谁？
A: FastGPT 的作者是 labring。
'''
原问题: Tell me about him
检索词: ["Introduce labring, the author of FastGPT." ," Background information on author labring." "," Why does labring do FastGPT?"]
----------------
历史记录:
'''
Q: 对话背景。
A: 关于 FatGPT 的介绍和使用等问题。
'''
原问题: 你好。
检索词: ["你好"]
----------------
历史记录:
'''
Q: FastGPT 如何收费？
A: FastGPT 收费可以参考……
'''
原问题: 你知道 laf 么？
检索词: ["laf 的官网地址是多少？","laf 的使用教程。","laf 有什么特点和优势。"]
----------------
历史记录:
'''
Q: FastGPT 的优势
A: 1. 开源
   2. 简便
   3. 扩展性强
'''
原问题: 介绍下第2点。
检索词: ["介绍下 FastGPT 简便的优势", "从哪些方面，可以体现出 FastGPT 的简便"]。
----------------
历史记录:
'''
Q: 什么是 FastGPT？
A: FastGPT 是一个 RAG 平台。
Q: 什么是 Laf？
A: Laf 是一个云函数开发平台。
'''
原问题: 它们有什么关系？
检索词: ["FastGPT和Laf有什么关系？","介绍下FastGPT","介绍下Laf"]
----------------
历史记录:
'''
{histories}
'''
原问题: {query}
检索词: """


class QAExt:
    llm = None

    def __init__(self, llm) -> None:
        self.llm = llm
        prompt = PromptTemplate.from_template(DEFAULT_PROMPT)
        # 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})