|
|
import asyncio |
|
|
import logging |
|
|
from typing import List, Optional, Sequence |
|
|
|
|
|
from langchain_core.callbacks import ( |
|
|
AsyncCallbackManagerForRetrieverRun, |
|
|
CallbackManagerForRetrieverRun, |
|
|
) |
|
|
from langchain_core.documents import Document |
|
|
from langchain_core.language_models import BaseLanguageModel |
|
|
from langchain_core.output_parsers import BaseOutputParser |
|
|
from langchain_core.prompts import BasePromptTemplate |
|
|
from langchain_core.prompts.prompt import PromptTemplate |
|
|
from langchain_core.retrievers import BaseRetriever |
|
|
from langchain_core.runnables import Runnable |
|
|
|
|
|
from langchain.chains.llm import LLMChain |
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
|
|
|
class LineListOutputParser(BaseOutputParser[List[str]]): |
|
|
"""Output parser for a list of lines.""" |
|
|
|
|
|
def parse(self, text: str) -> List[str]: |
|
|
lines = text.strip().split("\n") |
|
|
return list(filter(None, lines)) |
|
|
|
|
|
|
|
|
|
|
|
DEFAULT_QUERY_PROMPT = PromptTemplate( |
|
|
input_variables=["question"], |
|
|
template="""You are an AI language model assistant. Your task is |
|
|
to generate 3 different versions of the given user |
|
|
question to retrieve relevant documents from a vector database. |
|
|
By generating multiple perspectives on the user question, |
|
|
your goal is to help the user overcome some of the limitations |
|
|
of distance-based similarity search. Provide these alternative |
|
|
questions separated by newlines. Original question: {question}""", |
|
|
) |
|
|
|
|
|
|
|
|
def _unique_documents(documents: Sequence[Document]) -> List[Document]: |
|
|
return [doc for i, doc in enumerate(documents) if doc not in documents[:i]] |
|
|
|
|
|
|
|
|
class MultiQueryRetriever(BaseRetriever): |
|
|
"""Given a query, use an LLM to write a set of queries. |
|
|
|
|
|
Retrieve docs for each query. Return the unique union of all retrieved docs. |
|
|
""" |
|
|
|
|
|
retriever: BaseRetriever |
|
|
llm_chain: Runnable |
|
|
verbose: bool = True |
|
|
parser_key: str = "lines" |
|
|
"""DEPRECATED. parser_key is no longer used and should not be specified.""" |
|
|
include_original: bool = False |
|
|
"""Whether to include the original query in the list of generated queries.""" |
|
|
|
|
|
@classmethod |
|
|
def from_llm( |
|
|
cls, |
|
|
retriever: BaseRetriever, |
|
|
llm: BaseLanguageModel, |
|
|
prompt: BasePromptTemplate = DEFAULT_QUERY_PROMPT, |
|
|
parser_key: Optional[str] = None, |
|
|
include_original: bool = False, |
|
|
) -> "MultiQueryRetriever": |
|
|
"""Initialize from llm using default template. |
|
|
|
|
|
Args: |
|
|
retriever: retriever to query documents from |
|
|
llm: llm for query generation using DEFAULT_QUERY_PROMPT |
|
|
prompt: The prompt which aims to generate several different versions |
|
|
of the given user query |
|
|
include_original: Whether to include the original query in the list of |
|
|
generated queries. |
|
|
|
|
|
Returns: |
|
|
MultiQueryRetriever |
|
|
""" |
|
|
output_parser = LineListOutputParser() |
|
|
llm_chain = prompt | llm | output_parser |
|
|
return cls( |
|
|
retriever=retriever, |
|
|
llm_chain=llm_chain, |
|
|
include_original=include_original, |
|
|
) |
|
|
|
|
|
async def _aget_relevant_documents( |
|
|
self, |
|
|
query: str, |
|
|
*, |
|
|
run_manager: AsyncCallbackManagerForRetrieverRun, |
|
|
) -> List[Document]: |
|
|
"""Get relevant documents given a user query. |
|
|
|
|
|
Args: |
|
|
query: user query |
|
|
|
|
|
Returns: |
|
|
Unique union of relevant documents from all generated queries |
|
|
""" |
|
|
queries = await self.agenerate_queries(query, run_manager) |
|
|
if self.include_original: |
|
|
queries.append(query) |
|
|
documents = await self.aretrieve_documents(queries, run_manager) |
|
|
return self.unique_union(documents) |
|
|
|
|
|
async def agenerate_queries( |
|
|
self, question: str, run_manager: AsyncCallbackManagerForRetrieverRun |
|
|
) -> List[str]: |
|
|
"""Generate queries based upon user input. |
|
|
|
|
|
Args: |
|
|
question: user query |
|
|
|
|
|
Returns: |
|
|
List of LLM generated queries that are similar to the user input |
|
|
""" |
|
|
response = await self.llm_chain.ainvoke( |
|
|
{"question": question}, config={"callbacks": run_manager.get_child()} |
|
|
) |
|
|
if isinstance(self.llm_chain, LLMChain): |
|
|
lines = response["text"] |
|
|
else: |
|
|
lines = response |
|
|
if self.verbose: |
|
|
logger.info(f"Generated queries: {lines}") |
|
|
return lines |
|
|
|
|
|
async def aretrieve_documents( |
|
|
self, queries: List[str], run_manager: AsyncCallbackManagerForRetrieverRun |
|
|
) -> List[Document]: |
|
|
"""Run all LLM generated queries. |
|
|
|
|
|
Args: |
|
|
queries: query list |
|
|
|
|
|
Returns: |
|
|
List of retrieved Documents |
|
|
""" |
|
|
document_lists = await asyncio.gather( |
|
|
*( |
|
|
self.retriever.ainvoke( |
|
|
query, config={"callbacks": run_manager.get_child()} |
|
|
) |
|
|
for query in queries |
|
|
) |
|
|
) |
|
|
return [doc for docs in document_lists for doc in docs] |
|
|
|
|
|
def _get_relevant_documents( |
|
|
self, |
|
|
query: str, |
|
|
*, |
|
|
run_manager: CallbackManagerForRetrieverRun, |
|
|
) -> List[Document]: |
|
|
"""Get relevant documents given a user query. |
|
|
|
|
|
Args: |
|
|
query: user query |
|
|
|
|
|
Returns: |
|
|
Unique union of relevant documents from all generated queries |
|
|
""" |
|
|
queries = self.generate_queries(query, run_manager) |
|
|
if self.include_original: |
|
|
queries.append(query) |
|
|
documents = self.retrieve_documents(queries, run_manager) |
|
|
return self.unique_union(documents) |
|
|
|
|
|
def generate_queries( |
|
|
self, question: str, run_manager: CallbackManagerForRetrieverRun |
|
|
) -> List[str]: |
|
|
"""Generate queries based upon user input. |
|
|
|
|
|
Args: |
|
|
question: user query |
|
|
|
|
|
Returns: |
|
|
List of LLM generated queries that are similar to the user input |
|
|
""" |
|
|
response = self.llm_chain.invoke( |
|
|
{"question": question}, config={"callbacks": run_manager.get_child()} |
|
|
) |
|
|
if isinstance(self.llm_chain, LLMChain): |
|
|
lines = response["text"] |
|
|
else: |
|
|
lines = response |
|
|
if self.verbose: |
|
|
logger.info(f"Generated queries: {lines}") |
|
|
return lines |
|
|
|
|
|
def retrieve_documents( |
|
|
self, queries: List[str], run_manager: CallbackManagerForRetrieverRun |
|
|
) -> List[Document]: |
|
|
"""Run all LLM generated queries. |
|
|
|
|
|
Args: |
|
|
queries: query list |
|
|
|
|
|
Returns: |
|
|
List of retrieved Documents |
|
|
""" |
|
|
documents = [] |
|
|
for query in queries: |
|
|
docs = self.retriever.invoke( |
|
|
query, config={"callbacks": run_manager.get_child()} |
|
|
) |
|
|
documents.extend(docs) |
|
|
return documents |
|
|
|
|
|
def unique_union(self, documents: List[Document]) -> List[Document]: |
|
|
"""Get unique Documents. |
|
|
|
|
|
Args: |
|
|
documents: List of retrieved Documents |
|
|
|
|
|
Returns: |
|
|
List of unique retrieved Documents |
|
|
""" |
|
|
return _unique_documents(documents) |
|
|
|
