multi query

MultiQueryRetriever.py

@@ -0,0 +1,216 @@
+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.prompt import PromptTemplate
+from langchain_core.retrievers import BaseRetriever
+
+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 lines
+
+
+# Default prompt
+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]][:4]
+
+
+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: LLMChain
+    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: PromptTemplate = 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
+            include_original: Whether to include the original query in the list of
+                generated queries.
+
+        Returns:
+            MultiQueryRetriever
+        """
+        output_parser = LineListOutputParser()
+        llm_chain = LLMChain(llm=llm, prompt=prompt, output_parser=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:
+            question: 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.acall(
+            inputs={"question": question}, callbacks=run_manager.get_child()
+        )
+        lines = response["text"]
+        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.aget_relevant_documents(
+                    query, 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:
+            question: 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(
+            {"question": question}, callbacks=run_manager.get_child()
+        )
+        lines = response["text"]
+        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.get_relevant_documents(
+                query, callbacks=run_manager.get_child()
+            )
+            documents.extend(docs)
+        print("retrieve documents--", len(documents))  
+        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
+        """
+        print("unique union--", len(documents))
+        return _unique_documents(documents)

Retrieval.py

@@ -16,8 +16,7 @@ bsic_chain = RetrievalQA.from_chain_type(
 
 
 
-from langchain.retrievers.multi_query import MultiQueryRetriever
-# from kk import MultiQueryRetriever
+from MultiQueryRetriever import MultiQueryRetriever
 
 retriever_from_llm = MultiQueryRetriever.from_llm(
     retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),