bisheng-langchain 0.0.1__py3-none-any.whl

bisheng_langchain/chains/__init__.py

@@ -0,0 +1,5 @@
+from bisheng_langchain.chains.combine_documents.stuff import StuffDocumentsChain
+
+__all__ = [
+    'StuffDocumentsChain',
+]

bisheng_langchain/chains/combine_documents/stuff.py

@@ -0,0 +1,56 @@
+from typing import Any, List, Tuple
+
+from langchain.callbacks.manager import Callbacks
+from langchain.chains.combine_documents.stuff import StuffDocumentsChain as StuffDocumentsChainOld
+from langchain.docstore.document import Document
+
+
+class StuffDocumentsChain(StuffDocumentsChainOld):
+
+    token_max: int = -1
+
+    def combine_docs(self,
+                     docs: List[Document],
+                     callbacks: Callbacks = None,
+                     **kwargs: Any) -> Tuple[str, dict]:
+        """Stuff all documents into one prompt and pass to LLM.
+
+        Args:
+            docs: List of documents to join together into one variable
+            callbacks: Optional callbacks to pass along
+            **kwargs: additional parameters to use to get inputs to LLMChain.
+
+        Returns:
+            The first element returned is the single string output. The second
+            element returned is a dictionary of other keys to return.
+        """
+        inputs = self._get_inputs(docs, **kwargs)
+        # print('inputs:', len(inputs['context']))
+        # print('prompt_length:', self.prompt_length(docs, **kwargs))
+        if self.token_max > 0:
+            inputs[self.document_variable_name] = inputs[
+                self.document_variable_name][:self.token_max]
+        # Call predict on the LLM.
+        return self.llm_chain.predict(callbacks=callbacks, **inputs), {}
+
+    async def acombine_docs(self,
+                            docs: List[Document],
+                            callbacks: Callbacks = None,
+                            **kwargs: Any) -> Tuple[str, dict]:
+        """Stuff all documents into one prompt and pass to LLM.
+
+        Args:
+            docs: List of documents to join together into one variable
+            callbacks: Optional callbacks to pass along
+            **kwargs: additional parameters to use to get inputs to LLMChain.
+
+        Returns:
+            The first element returned is the single string output. The second
+            element returned is a dictionary of other keys to return.
+        """
+        inputs = self._get_inputs(docs, **kwargs)
+        if self.token_max > 0:
+            inputs[self.document_variable_name] = inputs[
+                self.document_variable_name][:self.token_max]
+        # Call predict on the LLM.
+        return await self.llm_chain.apredict(callbacks=callbacks, **inputs), {}

bisheng_langchain/chains/question_answering/__init__.py

@@ -0,0 +1,240 @@
+"""Load question answering chains."""
+from typing import Any, Mapping, Optional, Protocol
+
+# from langchain.chains.combine_documents.stuff import StuffDocumentsChain
+from bisheng_langchain.chains.combine_documents.stuff import StuffDocumentsChain
+from langchain.callbacks.base import BaseCallbackManager
+from langchain.callbacks.manager import Callbacks
+from langchain.chains import ReduceDocumentsChain
+from langchain.chains.combine_documents.base import BaseCombineDocumentsChain
+from langchain.chains.combine_documents.map_reduce import MapReduceDocumentsChain
+from langchain.chains.combine_documents.map_rerank import MapRerankDocumentsChain
+from langchain.chains.combine_documents.refine import RefineDocumentsChain
+from langchain.chains.llm import LLMChain
+from langchain.chains.question_answering import map_reduce_prompt, refine_prompts, stuff_prompt
+from langchain.chains.question_answering.map_rerank_prompt import PROMPT as MAP_RERANK_PROMPT
+from langchain.schema.language_model import BaseLanguageModel
+from langchain.schema.prompt_template import BasePromptTemplate
+
+
+class LoadingCallable(Protocol):
+    """Interface for loading the combine documents chain."""
+
+    def __call__(self, llm: BaseLanguageModel,
+                 **kwargs: Any) -> BaseCombineDocumentsChain:
+        """Callable to load the combine documents chain."""
+
+
+def _load_map_rerank_chain(
+    llm: BaseLanguageModel,
+    prompt: BasePromptTemplate = MAP_RERANK_PROMPT,
+    verbose: bool = False,
+    document_variable_name: str = 'context',
+    rank_key: str = 'score',
+    answer_key: str = 'answer',
+    callback_manager: Optional[BaseCallbackManager] = None,
+    callbacks: Callbacks = None,
+    **kwargs: Any,
+) -> MapRerankDocumentsChain:
+    llm_chain = LLMChain(
+        llm=llm,
+        prompt=prompt,
+        verbose=verbose,
+        callback_manager=callback_manager,
+        callbacks=callbacks,
+    )
+    return MapRerankDocumentsChain(
+        llm_chain=llm_chain,
+        rank_key=rank_key,
+        answer_key=answer_key,
+        document_variable_name=document_variable_name,
+        verbose=verbose,
+        callback_manager=callback_manager,
+        **kwargs,
+    )
+
+
+def _load_stuff_chain(
+    llm: BaseLanguageModel,
+    prompt: Optional[BasePromptTemplate] = None,
+    document_variable_name: str = 'context',
+    verbose: Optional[bool] = None,
+    callback_manager: Optional[BaseCallbackManager] = None,
+    callbacks: Callbacks = None,
+    **kwargs: Any,
+) -> StuffDocumentsChain:
+    _prompt = prompt or stuff_prompt.PROMPT_SELECTOR.get_prompt(llm)
+    llm_chain = LLMChain(
+        llm=llm,
+        prompt=_prompt,
+        verbose=verbose,
+        callback_manager=callback_manager,
+        callbacks=callbacks,
+    )
+    # TODO: document prompt
+    return StuffDocumentsChain(
+        llm_chain=llm_chain,
+        document_variable_name=document_variable_name,
+        verbose=verbose,
+        callback_manager=callback_manager,
+        **kwargs,
+    )
+
+
+def _load_map_reduce_chain(
+    llm: BaseLanguageModel,
+    question_prompt: Optional[BasePromptTemplate] = None,
+    combine_prompt: Optional[BasePromptTemplate] = None,
+    combine_document_variable_name: str = 'summaries',
+    map_reduce_document_variable_name: str = 'context',
+    collapse_prompt: Optional[BasePromptTemplate] = None,
+    reduce_llm: Optional[BaseLanguageModel] = None,
+    collapse_llm: Optional[BaseLanguageModel] = None,
+    verbose: Optional[bool] = None,
+    callback_manager: Optional[BaseCallbackManager] = None,
+    callbacks: Callbacks = None,
+    token_max: int = 3000,
+    **kwargs: Any,
+) -> MapReduceDocumentsChain:
+    _question_prompt = (
+        question_prompt
+        or map_reduce_prompt.QUESTION_PROMPT_SELECTOR.get_prompt(llm))
+    _combine_prompt = (
+        combine_prompt
+        or map_reduce_prompt.COMBINE_PROMPT_SELECTOR.get_prompt(llm))
+    map_chain = LLMChain(
+        llm=llm,
+        prompt=_question_prompt,
+        verbose=verbose,
+        callback_manager=callback_manager,
+        callbacks=callbacks,
+    )
+    _reduce_llm = reduce_llm or llm
+    reduce_chain = LLMChain(
+        llm=_reduce_llm,
+        prompt=_combine_prompt,
+        verbose=verbose,
+        callback_manager=callback_manager,
+        callbacks=callbacks,
+    )
+    # TODO: document prompt
+    combine_documents_chain = StuffDocumentsChain(
+        llm_chain=reduce_chain,
+        document_variable_name=combine_document_variable_name,
+        verbose=verbose,
+        callback_manager=callback_manager,
+        callbacks=callbacks,
+    )
+    if collapse_prompt is None:
+        collapse_chain = None
+        if collapse_llm is not None:
+            raise ValueError(
+                'collapse_llm provided, but collapse_prompt was not: please '
+                'provide one or stop providing collapse_llm.')
+    else:
+        _collapse_llm = collapse_llm or llm
+        collapse_chain = StuffDocumentsChain(
+            llm_chain=LLMChain(
+                llm=_collapse_llm,
+                prompt=collapse_prompt,
+                verbose=verbose,
+                callback_manager=callback_manager,
+                callbacks=callbacks,
+            ),
+            document_variable_name=combine_document_variable_name,
+            verbose=verbose,
+            callback_manager=callback_manager,
+        )
+    reduce_documents_chain = ReduceDocumentsChain(
+        combine_documents_chain=combine_documents_chain,
+        collapse_documents_chain=collapse_chain,
+        token_max=token_max,
+        verbose=verbose,
+    )
+    return MapReduceDocumentsChain(
+        llm_chain=map_chain,
+        document_variable_name=map_reduce_document_variable_name,
+        reduce_documents_chain=reduce_documents_chain,
+        verbose=verbose,
+        callback_manager=callback_manager,
+        callbacks=callbacks,
+        **kwargs,
+    )
+
+
+def _load_refine_chain(
+    llm: BaseLanguageModel,
+    question_prompt: Optional[BasePromptTemplate] = None,
+    refine_prompt: Optional[BasePromptTemplate] = None,
+    document_variable_name: str = 'context_str',
+    initial_response_name: str = 'existing_answer',
+    refine_llm: Optional[BaseLanguageModel] = None,
+    verbose: Optional[bool] = None,
+    callback_manager: Optional[BaseCallbackManager] = None,
+    callbacks: Callbacks = None,
+    **kwargs: Any,
+) -> RefineDocumentsChain:
+    _question_prompt = (
+        question_prompt
+        or refine_prompts.QUESTION_PROMPT_SELECTOR.get_prompt(llm))
+    _refine_prompt = refine_prompt or refine_prompts.REFINE_PROMPT_SELECTOR.get_prompt(
+        llm)
+    initial_chain = LLMChain(
+        llm=llm,
+        prompt=_question_prompt,
+        verbose=verbose,
+        callback_manager=callback_manager,
+        callbacks=callbacks,
+    )
+    _refine_llm = refine_llm or llm
+    refine_chain = LLMChain(
+        llm=_refine_llm,
+        prompt=_refine_prompt,
+        verbose=verbose,
+        callback_manager=callback_manager,
+        callbacks=callbacks,
+    )
+    return RefineDocumentsChain(
+        initial_llm_chain=initial_chain,
+        refine_llm_chain=refine_chain,
+        document_variable_name=document_variable_name,
+        initial_response_name=initial_response_name,
+        verbose=verbose,
+        callback_manager=callback_manager,
+        **kwargs,
+    )
+
+
+def load_qa_chain(
+    llm: BaseLanguageModel,
+    chain_type: str = 'stuff',
+    verbose: Optional[bool] = None,
+    callback_manager: Optional[BaseCallbackManager] = None,
+    **kwargs: Any,
+) -> BaseCombineDocumentsChain:
+    """Load question answering chain.
+
+    Args:
+        llm: Language Model to use in the chain.
+        chain_type: Type of document combining chain to use. Should be one of "stuff",
+            "map_reduce", "map_rerank", and "refine".
+        verbose: Whether chains should be run in verbose mode or not. Note that this
+            applies to all chains that make up the final chain.
+        callback_manager: Callback manager to use for the chain.
+
+    Returns:
+        A chain to use for question answering.
+    """
+    loader_mapping: Mapping[str, LoadingCallable] = {
+        'stuff': _load_stuff_chain,
+        'map_reduce': _load_map_reduce_chain,
+        'refine': _load_refine_chain,
+        'map_rerank': _load_map_rerank_chain,
+    }
+    if chain_type not in loader_mapping:
+        raise ValueError(f'Got unsupported chain type: {chain_type}. '
+                         f'Should be one of {loader_mapping.keys()}')
+    return loader_mapping[chain_type](llm,
+                                      verbose=verbose,
+                                      callback_manager=callback_manager,
+                                      **kwargs)

bisheng_langchain/chains/retrieval_qa/base.py

@@ -0,0 +1,89 @@
+from typing import List
+
+from langchain.callbacks.manager import AsyncCallbackManagerForChainRun, CallbackManagerForChainRun
+from langchain.chains.retrieval_qa.base import BaseRetrievalQA
+from langchain.pydantic_v1 import Field
+from langchain.schema import BaseRetriever, Document
+
+
+class MultiRetrievalQA(BaseRetrievalQA):
+    """Chain for question-answering against an index.
+
+    Example:
+        .. code-block:: python
+
+            from langchain.llms import OpenAI
+            from langchain.chains import RetrievalQA
+            from langchain.faiss import FAISS
+            from langchain.vectorstores.base import VectorStoreRetriever
+            retriever = VectorStoreRetriever(vectorstore=FAISS(...))
+            retrievalQA = RetrievalQA.from_llm(llm=OpenAI(), retriever=retriever)
+
+    """
+
+    vector_retriever: BaseRetriever = Field(exclude=True)
+    keyword_retriever: BaseRetriever = Field(exclude=True)
+    combine_strategy: str = 'keyword_front'  # "keyword_front, vector_front, mix"
+
+    def _get_docs(
+        self,
+        question: str,
+        *,
+        run_manager: CallbackManagerForChainRun,
+    ) -> List[Document]:
+        """Get docs."""
+        vector_docs = self.vector_retriever.get_relevant_documents(
+            question, callbacks=run_manager.get_child())
+        keyword_docs = self.keyword_retriever.get_relevant_documents(
+            question, callbacks=run_manager.get_child())
+        if self.combine_strategy == 'keyword_front':
+            return keyword_docs + vector_docs
+        elif self.combine_strategy == 'vector_front':
+            return vector_docs + keyword_docs
+        elif self.combine_strategy == 'mix':
+            combine_docs = []
+            min_len = min(len(keyword_docs), len(vector_docs))
+            for i in range(min_len):
+                combine_docs.append(keyword_docs[i])
+                combine_docs.append(vector_docs[i])
+            combine_docs.extend(keyword_docs[min_len:])
+            combine_docs.extend(vector_docs[min_len:])
+            return combine_docs
+        else:
+            raise ValueError(f'Expected combine_strategy to be one of '
+                             f'(keyword_front, vector_front, mix),'
+                             f'instead found {self.combine_strategy}')
+
+    async def _aget_docs(
+        self,
+        question: str,
+        *,
+        run_manager: AsyncCallbackManagerForChainRun,
+    ) -> List[Document]:
+        """Get docs."""
+        vector_docs = await self.vector_retriever.get_relevant_documents(
+            question, callbacks=run_manager.get_child())
+        keyword_docs = await self.keyword_retriever.get_relevant_documents(
+            question, callbacks=run_manager.get_child())
+        if self.combine_strategy == 'keyword_front':
+            return keyword_docs + vector_docs
+        elif self.combine_strategy == 'vector_front':
+            return vector_docs + keyword_docs
+        elif self.combine_strategy == 'mix':
+            combine_docs = []
+            min_len = min(len(keyword_docs), len(vector_docs))
+            for i in range(min_len):
+                combine_docs.append(keyword_docs[i])
+                combine_docs.append(vector_docs[i])
+            combine_docs.extend(keyword_docs[min_len:])
+            combine_docs.extend(vector_docs[min_len:])
+            return combine_docs
+        else:
+            raise ValueError(f'Expected combine_strategy to be one of '
+                             f'(keyword_front, vector_front, mix),'
+                             f'instead found {self.combine_strategy}')
+
+    @property
+    def _chain_type(self) -> str:
+        """Return the chain type."""
+        return 'multi_retrieval_qa'

bisheng_langchain/chat_models/__init__.py

@@ -0,0 +1,11 @@
+from .host_llm import BaichuanChat, ChatGLM2Host, Llama2Chat, QwenChat
+from .minimax import ChatMinimaxAI
+from .proxy_llm import ProxyChatLLM
+from .wenxin import ChatWenxin
+from .xunfeiai import ChatXunfeiAI
+from .zhipuai import ChatZhipuAI
+
+__all__ = [
+    'ProxyChatLLM', 'ChatMinimaxAI', 'ChatWenxin', 'ChatZhipuAI',
+    'ChatXunfeiAI', 'Llama2Chat', 'ChatGLM2Host', 'BaichuanChat', 'QwenChat'
+]