AI多跳RAG折腾手记

from langchain.retrievers.multi_hop import MultiHopRetriever
from langchain_community.embeddings import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma

# 基础向量库
vectorstore = Chroma(
    embedding_function=OpenAIEmbeddings(),
    persist_directory="./chroma_db"
)

# 创建多跳检索器
retriever = MultiHopRetriever(
    retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
    max_hops=3,
    verbose=True
)

# 使用
question = "项目A和项目B在架构设计上有什么不同？"
docs = retriever.get_relevant_documents(question)

from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate

class ChainHopRAG:
    def __init__(self, vectorstore, max_hops=3):
        self.vectorstore = vectorstore
        self.max_hops = max_hops
        self.llm = ChatOpenAI(temperature=0)

        self.query_gen_prompt = ChatPromptTemplate.from_messages([
            ("system", "你是一个专业的查询生成助手。根据已收集的信息，判断是否需要进一步检索，如果需要，生成最合适的检索查询。"),
            ("user", "原始问题：{question}\n已收集信息：{collected_info}\n\n请判断：1. 是否需要进一步检索？2. 如果需要，生成1-2个检索查询。")
        ])

        self.answer_prompt = ChatPromptTemplate.from_messages([
            ("system", "你是一个专业的问答助手。基于收集到的所有信息回答用户问题。"),
            ("user", "问题：{question}\n收集到的信息：\n{collected_info}")
        ])

    def run(self, question):
        collected_info = []
        current_queries = [question]

        for hop in range(self.max_hops):
            # 检索当前查询
            for query in current_queries:
                docs = self.vectorstore.similarity_search(query, k=3)
                collected_info.extend(docs)

            # 判断是否需要继续
            if hop < self.max_hops - 1:
                response = self.llm.invoke(
                    self.query_gen_prompt.format(
                        question=question,
                        collected_info="\n".join([d.page_content for d in collected_info[-6:]])
                    )
                )

                # 解析模型反馈
                if "不需要进一步检索" in response.content:
                    break

                # 提取新查询
                current_queries = self._extract_queries(response.content)
                if not current_queries:
                    break
            else:
                break

        # 生成最终答案
        final_answer = self.llm.invoke(
            self.answer_prompt.format(
                question=question,
                collected_info="\n".join([d.page_content for d in collected_info])
            )
        )

        return final_answer.content

    def _extract_queries(self, response_text):
        # 从响应中提取查询语句
        lines = response_text.split('\n')
        queries = []
        for line in lines:
            if line.strip().startswith('-') or line.strip().startswith('•'):
                queries.append(line.strip()[1:].strip())
        return queries

def filter_relevant_docs(docs, question, threshold=0.7):
    """过滤相关性较低的文档"""
    relevant_docs = []
    for doc in docs:
        relevance = calculate_relevance(doc.page_content, question)
        if relevance >= threshold:
            relevant_docs.append(doc)
    return relevant_docs

QUERY_TEMPLATES = {
    "comparison": "对比{entity1}和{entity2}在{aspect}方面的差异",
    "reasoning": "为什么{entity}采用了{approach}，基于哪些考虑",
    "timeline": "{entity}在{time_period}的发展历程和关键节点"
}

def should_use_multi_hop(question):
    """判断问题是否适合多跳RAG"""
    keywords = ["对比", "差异", "为什么", "关系", "影响"]
    return any(keyword in question for keyword in keywords)