retriever.py

import argparse
import configparser
import os
import requests
import time
import uvicorn

from fastapi import FastAPI, Request
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from sklearn.metrics import precision_recall_curve
from loguru import logger
from BCEmbedding.tools.langchain import BCERerank
from langchain.retrievers import ContextualCompressionRetriever
from langchain_community.vectorstores.utils import DistanceStrategy
from requests.exceptions import RequestException
from elastic_keywords_search import ElasticKeywordsSearch

app = FastAPI()


class Retriever:
    def __init__(self, config) -> None:

        self.mix, self.es, self.vector = None, None, None
        work_dir = config['default']['work_dir']

        self.es_top_k = int(config['rag']['es_top_k'])
        self.vector_top_k = int(config['rag']['vector_top_k'])

        embedding_model_path = config.get('rag', 'embedding_model_path') or None
        reranker_model_path = config.get('rag', 'reranker_model_path') or None
        es_url = config.get('rag', 'es_url') or None
        index_name = config.get('rag', 'index_name') or None

        # Mix
        if embedding_model_path and reranker_model_path and es_url and index_name:
            self.init_mix_retriever(work_dir, embedding_model_path, reranker_model_path, es_url, index_name)

        # ES
        elif not embedding_model_path or not reranker_model_path:
            if self.is_es_available(es_url, index_name):
                self.es_retriever = ElasticKeywordsSearch(es_url, index_name, drop_old=False)
                self.weights = [0.5, 0.5]
                self.es = True
                logger.info('Initializing ES retriever alone!')

        # Vector
        elif not es_url or not index_name:
            self.init_vector_retriever(work_dir, embedding_model_path, reranker_model_path)
            self.vector = True
            logger.info('Initializing Vector retriever alone!')
        else:
            raise ValueError(
                "Incomplete configuration. Please specify all required parameters for either vector or ES retrieval.")

    def init_vector_retriever(self, work_dir, embedding_model_path, reranker_model_path):
        logger.info('loading test2vec and rerank models')
        self.embeddings = HuggingFaceEmbeddings(
            model_name=embedding_model_path,
            model_kwargs={'device': 'cuda'},
            encode_kwargs={
                'batch_size': 1,
                'normalize_embeddings': True
            })
        # half
        self.embeddings.client = self.embeddings.client.half()
        reranker_args = {
            'model': reranker_model_path,
            'top_n': self.vector_top_k,
            'device': 'cuda',
            'use_fp16': True
        }
        self.reranker = BCERerank(**reranker_args)

        self.vector_store = FAISS.load_local(
            os.path.join(work_dir, 'db_response'),
            embeddings=self.embeddings,
            allow_dangerous_deserialization=True,
            distance_strategy=DistanceStrategy.MAX_INNER_PRODUCT)

        retriever = self.vector_store.as_retriever(
            search_type='similarity',
            search_kwargs={
                'score_threshold': 0.15,
                'k': 30
            }
        )
        self.compression_retriever = ContextualCompressionRetriever(
            base_compressor=self.reranker, base_retriever=retriever)

    def init_mix_retriever(self, work_dir, embedding_model_path, reranker_model_path, es_url, index_name):

        if self.is_es_available(es_url, index_name):
            self.es_retriever = ElasticKeywordsSearch(es_url, index_name, drop_old=False)
            self.weights = [0.5, 0.5]
            self.init_vector_retriever(work_dir, embedding_model_path, reranker_model_path)
            self.mix = True
            logger.info('Initializing Mix retriever!')
        else:
            self.init_vector_retriever(work_dir, embedding_model_path, reranker_model_path)
            self.vector = True
            logger.info('Initializing Vector retriever alone!')

    def is_es_available(self, url, index_name, timeout=5):
        try:
            response = requests.get(f"{url}/_cluster/health", timeout=timeout)
            if response.status_code == 200:

                index_response = requests.head(f"{url}/{index_name}", timeout=timeout)
                if index_response.status_code == 200:
                    logger.info(f"The index:'{index_name}' exist!")
                    return True
                elif index_response.status_code == 404:
                    logger.warning(f"The index:'{index_name}' not exist!")
                else:
                    logger.error(f"Unexpected status code when checking index: {index_response.status_code}")
            else:
                logger.error(f"Elasticsearch service returned non-200 status code: {response.status_code}")
        except RequestException as e:
            logger.error(f"Error connecting to Elasticsearch service: {e}")

        return False

    def weighted_reciprocal_rank(self, es_docs, vector_docs):
        # Create a union of all unique documents in the input doc_lists
        all_documents = set()

        for vector_doc in vector_docs:
            all_documents.add(vector_doc.page_content)
        for es_doc in es_docs:
            all_documents.add(es_doc.page_content)

        rrf_score_dic = {doc: 0.0 for doc in all_documents}

        for rank, vector_doc in enumerate(vector_docs, start=1):
            rrf_score = self.weights[1] * (1 / (rank + 60))
            rrf_score_dic[vector_doc.page_content] += rrf_score
        for rank, es_doc in enumerate(es_docs, start=1):
            rrf_score = self.weights[0] * (1 / (rank + 60))
            rrf_score_dic[es_doc.page_content] += rrf_score

        sorted_documents = sorted(rrf_score_dic.keys(), key=lambda x: rrf_score_dic[x], reverse=True)

        # Map the sorted page_content back to the original document objects
        page_content_to_doc_map = {}
        for doc in es_docs:
            page_content_to_doc_map[doc.page_content] = doc
        for doc in vector_docs:
            page_content_to_doc_map[doc.page_content] = doc

        sorted_docs = [page_content_to_doc_map[page_content] for page_content in sorted_documents]

        return sorted_docs

    def remove_duplicates(self, sorted_docs):
        seen = set()
        unique_docs = []

        for doc in sorted_docs:
            identifier = (
                doc.metadata.get('source', ''),
                doc.metadata.get('read', ''),
                # doc.page_content  # Need further testing
            )

            if identifier not in seen:
                seen.add(identifier)
                unique_docs.append(doc)

        return unique_docs

    def hybrid_retrieval(self, query):
        es_docs = self.es_retriever.similarity_search_with_score(query, k=self.es_top_k)
        vector_docs = self.query(query)
        sorted_docs = self.weighted_reciprocal_rank(es_docs, vector_docs)
        unique_docs = self.remove_duplicates(sorted_docs)

        return unique_docs

    def rag_workflow(self, query):
        chunks = []
        time_1 = time.time()

        if self.mix:
            chunks = self.hybrid_retrieval(query)
        elif self.es:
            chunks = self.es_retriever.similarity_search_with_score(query, k=self.es_top_k)
        else:
            chunks = self.query(query)

        time_2 = time.time()
        logger.debug(f'query:{query} \nchunks:{chunks} \ntimecost:{time_2 - time_1}')
        return chunks

    def query(self, question: str):
        if question is None or len(question) < 1:
            return None

        if len(question) > 512:
            logger.warning('input too long, truncate to 512')
            question = question[0:512]

        docs = self.compression_retriever.get_relevant_documents(question)
        return docs


retriever = None


@app.post("/retrieve")
async def retrieve(request: Request):
    data = await request.json()
    query = data.get("query")
    chunks = retriever.rag_workflow(query)
    return {"chunks": chunks}


def rag_retrieve(args: str):
    """
    启动 Web 服务器，接收 HTTP 请求，并通过调用本地的 rag 检索服务.

    """
    global retriever
    config = configparser.ConfigParser()
    config.read(args.config_path)
    bind_port = int(config['default']['bind_port'])
    os.environ["CUDA_VISIBLE_DEVICES"] = args.dcu_id
    retriever = Retriever(config)

    uvicorn.run(app, host="0.0.0.0", port=bind_port)


def parse_args():
    parser = argparse.ArgumentParser(
        description='Feature store for processing directories.')
    parser.add_argument(
        '--config_path',
        default='/ai/rag/config.ini',
        help='config目录')
    parser.add_argument(
        '--dcu_id',
        default=None,
        help='设置DCU')
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    rag_retrieve(args)


if __name__ == '__main__':
    main()