43 lines
1.8 KiB
Python
43 lines
1.8 KiB
Python
|
# rag/query_transform.py
|
|||
|
|
|
|||
|
|
from typing import List
|
|||
|
|
from langchain_core.language_models import BaseLanguageModel
|
|||
|
|
from langchain_core.prompts import PromptTemplate
|
|||
|
|
|
|||
|
|
MULTI_QUERY_PROMPT = PromptTemplate.from_template(
|
|||
|
|
"""你是一个专业的查询改写助手。你的任务是将用户的问题改写成 {num_queries} 个不同的版本。
|
|||
|
|
这些版本应该从不同的角度、使用不同的关键词来表达相同或相关的意图。
|
|||
|
|
|
|||
|
|
原始问题: {question}
|
|||
|
|
|
|||
|
|
请生成 {num_queries} 个不同版本的查询，每个版本一行。
|
|||
|
|
确保每个版本都是独立、完整的查询语句。
|
|||
|
|
|
|||
|
|
生成 {num_queries} 个查询:"""
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
class MultiQueryGenerator:
|
|||
|
|
"""多路查询生成器（不依赖 LangChain 的 MultiQueryRetriever）"""
|
|||
|
|
|
|||
|
|
def __init__(self, llm: BaseLanguageModel, num_queries: int = 3):
|
|||
|
|
self.llm = llm
|
|||
|
|
self.num_queries = num_queries
|
|||
|
|
self.prompt = MULTI_QUERY_PROMPT
|
|||
|
|
|
|||
|
|
def generate(self, query: str) -> List[str]:
|
|||
|
|
"""同步生成多个查询变体"""
|
|||
|
|
prompt_str = self.prompt.format(num_queries=self.num_queries, question=query)
|
|||
|
|
response = self.llm.invoke(prompt_str)
|
|||
|
|
# 处理响应内容，按行分割并去除空行和首尾空白
|
|||
|
|
lines = response.content.strip().split('\n')
|
|||
|
|
queries = [line.strip() for line in lines if line.strip()]
|
|||
|
|
# 确保至少返回原始查询
|
|||
|
|
return queries[:self.num_queries] if queries else [query]
|
|||
|
|
|
|||
|
|
async def agenerate(self, query: str) -> List[str]:
|
|||
|
|
"""异步生成多个查询变体"""
|
|||
|
|
prompt_str = self.prompt.format(num_queries=self.num_queries, question=query)
|
|||
|
|
response = await self.llm.ainvoke(prompt_str)
|
|||
|
|
lines = response.content.strip().split('\n')
|
|||
|
|
queries = [line.strip() for line in lines if line.strip()]
|
|||
|
|
return queries[:self.num_queries] if queries else [query]
|